1# Copyright 2014 The Chromium Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5from __future__ import division
6
7import warnings
8
9from telemetry.internal.util import external_modules
10from telemetry.util import color_histogram
11from telemetry.util import rgba_color
12import png
13
14cv2 = external_modules.ImportOptionalModule('cv2')
15np = external_modules.ImportRequiredModule('numpy')
16
17
18def Channels(image):
19  return image.shape[2]
20
21def Width(image):
22  return image.shape[1]
23
24def Height(image):
25  return image.shape[0]
26
27def Pixels(image):
28  return bytearray(np.uint8(image[:, :, ::-1]).flat)  # Convert from bgr to rgb.
29
30def GetPixelColor(image, x, y):
31  bgr = image[y][x]
32  return rgba_color.RgbaColor(bgr[2], bgr[1], bgr[0])
33
34def WritePngFile(image, path):
35  if cv2 is not None:
36    cv2.imwrite(path, image)
37  else:
38    with open(path, "wb") as f:
39      metadata = {}
40      metadata['size'] = (Width(image), Height(image))
41      metadata['alpha'] = False
42      metadata['bitdepth'] = 8
43      img = image[:, :, ::-1]
44      pixels = img.reshape(-1).tolist()
45      png.Writer(**metadata).write_array(f, pixels)
46
47def FromRGBPixels(width, height, pixels, bpp):
48  img = np.array(pixels, order='F', dtype=np.uint8)
49  img.resize((height, width, bpp))
50  if bpp == 4:
51    img = img[:, :, :3]  # Drop alpha.
52  return img[:, :, ::-1]  # Convert from rgb to bgr.
53
54def FromPngFile(path):
55  if cv2 is not None:
56    img = cv2.imread(path, cv2.CV_LOAD_IMAGE_COLOR)
57    if img is None:
58      raise ValueError('Image at path {0} could not be read'.format(path))
59    return img
60  else:
61    with open(path, "rb") as f:
62      return FromPng(f.read())
63
64def FromPng(png_data):
65  if cv2 is not None:
66    file_bytes = np.asarray(bytearray(png_data), dtype=np.uint8)
67    return cv2.imdecode(file_bytes, cv2.CV_LOAD_IMAGE_COLOR)
68  else:
69    warnings.warn(
70        'Using pure python png decoder, which could be very slow. To speed up, '
71        'consider installing numpy & cv2 (OpenCV).')
72    width, height, pixels, meta = png.Reader(bytes=png_data).read_flat()
73    return FromRGBPixels(width, height, pixels, 4 if meta['alpha'] else 3)
74
75def _SimpleDiff(image1, image2):
76  if cv2 is not None:
77    return cv2.absdiff(image1, image2)
78  else:
79    amax = np.maximum(image1, image2)
80    amin = np.minimum(image1, image2)
81    return amax - amin
82
83def AreEqual(image1, image2, tolerance, likely_equal):
84  if image1.shape != image2.shape:
85    return False
86  self_image = image1
87  other_image = image2
88  if tolerance:
89    if likely_equal:
90      return np.amax(_SimpleDiff(image1, image2)) <= tolerance
91    else:
92      for row in xrange(Height(image1)):
93        if np.amax(_SimpleDiff(image1[row], image2[row])) > tolerance:
94          return False
95      return True
96  else:
97    if likely_equal:
98      return (self_image == other_image).all()
99    else:
100      for row in xrange(Height(image1)):
101        if not (self_image[row] == other_image[row]).all():
102          return False
103      return True
104
105def Diff(image1, image2):
106  self_image = image1
107  other_image = image2
108  if image1.shape[2] != image2.shape[2]:
109    raise ValueError('Cannot diff images of differing bit depth')
110  if image1.shape[:2] != image2.shape[:2]:
111    width = max(Width(image1), Width(image2))
112    height = max(Height(image1), Height(image2))
113    self_image = np.zeros((width, height, image1.shape[2]), np.uint8)
114    other_image = np.zeros((width, height, image1.shape[2]), np.uint8)
115    self_image[0:Height(image1), 0:Width(image1)] = image1
116    other_image[0:Height(image2), 0:Width(image2)] = image2
117  return _SimpleDiff(self_image, other_image)
118
119def GetBoundingBox(image, color, tolerance):
120  if cv2 is not None:
121    color = np.array([color.b, color.g, color.r])
122    img = cv2.inRange(image, np.subtract(color[0:3], tolerance),
123                      np.add(color[0:3], tolerance))
124    count = cv2.countNonZero(img)
125    if count == 0:
126      return None, 0
127    contours, _ = cv2.findContours(img, cv2.RETR_LIST, cv2.CHAIN_APPROX_NONE)
128    contour = np.concatenate(contours)
129    return cv2.boundingRect(contour), count
130  else:
131    if tolerance:
132      color = np.array([color.b, color.g, color.r])
133      colorm = color - tolerance
134      colorp = color + tolerance
135      b = image[:, :, 0]
136      g = image[:, :, 1]
137      r = image[:, :, 2]
138      w = np.where(((b >= colorm[0]) & (b <= colorp[0]) &
139                    (g >= colorm[1]) & (g <= colorp[1]) &
140                    (r >= colorm[2]) & (r <= colorp[2])))
141    else:
142      w = np.where((image[:, :, 0] == color.b) &
143                   (image[:, :, 1] == color.g) &
144                   (image[:, :, 2] == color.r))
145    if len(w[0]) == 0:
146      return None, 0
147    return (w[1][0], w[0][0], w[1][-1] - w[1][0] + 1, w[0][-1] - w[0][0] + 1), \
148        len(w[0])
149
150def Crop(image, left, top, width, height):
151  img_height, img_width = image.shape[:2]
152  if (left < 0 or top < 0 or
153      (left + width) > img_width or
154      (top + height) > img_height):
155    raise ValueError('Invalid dimensions')
156  return image[top:top + height, left:left + width]
157
158def GetColorHistogram(image, ignore_color, tolerance):
159  if cv2 is not None:
160    mask = None
161    if ignore_color is not None:
162      color = np.array([ignore_color.b, ignore_color.g, ignore_color.r])
163      mask = ~cv2.inRange(image, np.subtract(color, tolerance),
164                          np.add(color, tolerance))
165
166    flatten = np.ndarray.flatten
167    hist_b = flatten(cv2.calcHist([image], [0], mask, [256], [0, 256]))
168    hist_g = flatten(cv2.calcHist([image], [1], mask, [256], [0, 256]))
169    hist_r = flatten(cv2.calcHist([image], [2], mask, [256], [0, 256]))
170  else:
171    filtered = image.reshape(-1, 3)
172    if ignore_color is not None:
173      color = np.array([ignore_color.b, ignore_color.g, ignore_color.r])
174      colorm = np.array(color) - tolerance
175      colorp = np.array(color) + tolerance
176      in_range = ((filtered[:, 0] < colorm[0]) | (filtered[:, 0] > colorp[0]) |
177                  (filtered[:, 1] < colorm[1]) | (filtered[:, 1] > colorp[1]) |
178                  (filtered[:, 2] < colorm[2]) | (filtered[:, 2] > colorp[2]))
179      filtered = np.compress(in_range, filtered, axis=0)
180    if len(filtered[:, 0]) == 0:
181      return color_histogram.ColorHistogram(np.zeros((256)), np.zeros((256)),
182                                      np.zeros((256)), ignore_color)
183    hist_b = np.bincount(filtered[:, 0], minlength=256)
184    hist_g = np.bincount(filtered[:, 1], minlength=256)
185    hist_r = np.bincount(filtered[:, 2], minlength=256)
186
187  return color_histogram.ColorHistogram(hist_r, hist_g, hist_b, ignore_color)
188