1"""
2Utility functions to deal with ppm (qemu screendump format) files.
3
4@copyright: Red Hat 2008-2009
5"""
6
7import os, struct, time, re
8from autotest_lib.client.bin import utils
9
10# Some directory/filename utils, for consistency
11
12def find_id_for_screendump(md5sum, dir):
13    """
14    Search dir for a PPM file whose name ends with md5sum.
15
16    @param md5sum: md5 sum string
17    @param dir: Directory that holds the PPM files.
18    @return: The file's basename without any preceding path, e.g.
19    '20080101_120000_d41d8cd98f00b204e9800998ecf8427e.ppm'.
20    """
21    try:
22        files = os.listdir(dir)
23    except OSError:
24        files = []
25    for file in files:
26        exp = re.compile(r"(.*_)?" + md5sum + r"\.ppm", re.IGNORECASE)
27        if exp.match(file):
28            return file
29
30
31def generate_id_for_screendump(md5sum, dir):
32    """
33    Generate a unique filename using the given MD5 sum.
34
35    @return: Only the file basename, without any preceding path. The
36    filename consists of the current date and time, the MD5 sum and a .ppm
37    extension, e.g. '20080101_120000_d41d8cd98f00b204e9800998ecf8427e.ppm'.
38    """
39    filename = time.strftime("%Y%m%d_%H%M%S") + "_" + md5sum + ".ppm"
40    return filename
41
42
43def get_data_dir(steps_filename):
44    """
45    Return the data dir of the given steps filename.
46    """
47    filename = os.path.basename(steps_filename)
48    return os.path.join(os.path.dirname(steps_filename), "..", "steps_data",
49                        filename + "_data")
50
51
52# Functions for working with PPM files
53
54def image_read_from_ppm_file(filename):
55    """
56    Read a PPM image.
57
58    @return: A 3 element tuple containing the width, height and data of the
59            image.
60    """
61    fin = open(filename,"rb")
62    l1 = fin.readline()
63    l2 = fin.readline()
64    l3 = fin.readline()
65    data = fin.read()
66    fin.close()
67
68    (w, h) = map(int, l2.split())
69    return (w, h, data)
70
71
72def image_write_to_ppm_file(filename, width, height, data):
73    """
74    Write a PPM image with the given width, height and data.
75
76    @param filename: PPM file path
77    @param width: PPM file width (pixels)
78    @param height: PPM file height (pixels)
79    """
80    fout = open(filename,"wb")
81    fout.write("P6\n")
82    fout.write("%d %d\n" % (width, height))
83    fout.write("255\n")
84    fout.write(data)
85    fout.close()
86
87
88def image_crop(width, height, data, x1, y1, dx, dy):
89    """
90    Crop an image.
91
92    @param width: Original image width
93    @param height: Original image height
94    @param data: Image data
95    @param x1: Desired x coordinate of the cropped region
96    @param y1: Desired y coordinate of the cropped region
97    @param dx: Desired width of the cropped region
98    @param dy: Desired height of the cropped region
99    @return: A 3-tuple containing the width, height and data of the
100    cropped image.
101    """
102    if x1 > width - 1: x1 = width - 1
103    if y1 > height - 1: y1 = height - 1
104    if dx > width - x1: dx = width - x1
105    if dy > height - y1: dy = height - y1
106    newdata = ""
107    index = (x1 + y1*width) * 3
108    for i in range(dy):
109        newdata += data[index:(index+dx*3)]
110        index += width*3
111    return (dx, dy, newdata)
112
113
114def image_md5sum(width, height, data):
115    """
116    Return the md5sum of an image.
117
118    @param width: PPM file width
119    @param height: PPM file height
120    @data: PPM file data
121    """
122    header = "P6\n%d %d\n255\n" % (width, height)
123    hash = utils.hash('md5', header)
124    hash.update(data)
125    return hash.hexdigest()
126
127
128def get_region_md5sum(width, height, data, x1, y1, dx, dy,
129                      cropped_image_filename=None):
130    """
131    Return the md5sum of a cropped region.
132
133    @param width: Original image width
134    @param height: Original image height
135    @param data: Image data
136    @param x1: Desired x coord of the cropped region
137    @param y1: Desired y coord of the cropped region
138    @param dx: Desired width of the cropped region
139    @param dy: Desired height of the cropped region
140    @param cropped_image_filename: if not None, write the resulting cropped
141            image to a file with this name
142    """
143    (cw, ch, cdata) = image_crop(width, height, data, x1, y1, dx, dy)
144    # Write cropped image for debugging
145    if cropped_image_filename:
146        image_write_to_ppm_file(cropped_image_filename, cw, ch, cdata)
147    return image_md5sum(cw, ch, cdata)
148
149
150def image_verify_ppm_file(filename):
151    """
152    Verify the validity of a PPM file.
153
154    @param filename: Path of the file being verified.
155    @return: True if filename is a valid PPM image file. This function
156    reads only the first few bytes of the file so it should be rather fast.
157    """
158    try:
159        size = os.path.getsize(filename)
160        fin = open(filename, "rb")
161        assert(fin.readline().strip() == "P6")
162        (width, height) = map(int, fin.readline().split())
163        assert(width > 0 and height > 0)
164        assert(fin.readline().strip() == "255")
165        size_read = fin.tell()
166        fin.close()
167        assert(size - size_read == width*height*3)
168        return True
169    except:
170        return False
171
172
173def image_comparison(width, height, data1, data2):
174    """
175    Generate a green-red comparison image from two given images.
176
177    @param width: Width of both images
178    @param height: Height of both images
179    @param data1: Data of first image
180    @param data2: Data of second image
181    @return: A 3-element tuple containing the width, height and data of the
182            generated comparison image.
183
184    @note: Input images must be the same size.
185    """
186    newdata = ""
187    i = 0
188    while i < width*height*3:
189        # Compute monochromatic value of current pixel in data1
190        pixel1_str = data1[i:i+3]
191        temp = struct.unpack("BBB", pixel1_str)
192        value1 = int((temp[0] + temp[1] + temp[2]) / 3)
193        # Compute monochromatic value of current pixel in data2
194        pixel2_str = data2[i:i+3]
195        temp = struct.unpack("BBB", pixel2_str)
196        value2 = int((temp[0] + temp[1] + temp[2]) / 3)
197        # Compute average of the two values
198        value = int((value1 + value2) / 2)
199        # Scale value to the upper half of the range [0, 255]
200        value = 128 + value / 2
201        # Compare pixels
202        if pixel1_str == pixel2_str:
203            # Equal -- give the pixel a greenish hue
204            newpixel = [0, value, 0]
205        else:
206            # Not equal -- give the pixel a reddish hue
207            newpixel = [value, 0, 0]
208        newdata += struct.pack("BBB", newpixel[0], newpixel[1], newpixel[2])
209        i += 3
210    return (width, height, newdata)
211
212
213def image_fuzzy_compare(width, height, data1, data2):
214    """
215    Return the degree of equality of two given images.
216
217    @param width: Width of both images
218    @param height: Height of both images
219    @param data1: Data of first image
220    @param data2: Data of second image
221    @return: Ratio equal_pixel_count / total_pixel_count.
222
223    @note: Input images must be the same size.
224    """
225    equal = 0.0
226    different = 0.0
227    i = 0
228    while i < width*height*3:
229        pixel1_str = data1[i:i+3]
230        pixel2_str = data2[i:i+3]
231        # Compare pixels
232        if pixel1_str == pixel2_str:
233            equal += 1.0
234        else:
235            different += 1.0
236        i += 3
237    return equal / (equal + different)
238