1# Copyright 2024 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#      http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Utility functions for low light camera tests."""
15
16import logging
17import os.path
18
19import cv2
20import matplotlib.pyplot as plt
21import numpy as np
22
23_LOW_LIGHT_BOOST_AVG_DELTA_LUMINANCE_THRESH = 18
24_LOW_LIGHT_BOOST_AVG_LUMINANCE_THRESH = 90
25_BOUNDING_BOX_COLOR = (0, 255, 0)
26_BOX_MIN_SIZE_RATIO = 0.08  # 8% of the cropped image width
27_BOX_MAX_SIZE_RATIO = 0.5  # 50% of the cropped image width
28_BOX_PADDING_RATIO = 0.2
29_CROP_PADDING = 10
30_EXPECTED_NUM_OF_BOXES = 20  # The captured image must result in 20 detected
31                             # boxes since the test scene has 20 boxes
32_KEY_BOTTOM_LEFT = 'bottom_left'
33_KEY_BOTTOM_RIGHT = 'bottom_right'
34_KEY_TOP_LEFT = 'top_left'
35_KEY_TOP_RIGHT = 'top_right'
36_MAX_ASPECT_RATIO = 1.2
37_MIN_ASPECT_RATIO = 0.8
38_RED_BGR_COLOR = (0, 0, 255)
39_NUM_CLUSTERS = 8
40_K_MEANS_ITERATIONS = 10
41_K_MEANS_EPSILON = 0.5
42_TEXT_COLOR = (255, 255, 255)
43
44# pylint: disable=line-too-long
45# Allowed tablets for low light scenes
46# List entries must be entered in lowercase
47TABLET_LOW_LIGHT_SCENES_ALLOWLIST = (
48    'gta8wifi',  # Samsung Galaxy Tab A8
49    'gta8',  # Samsung Galaxy Tab A8 LTE
50    'gta9pwifi',  # Samsung Galaxy Tab A9+
51)
52
53
54def _crop(img):
55  """Crops the captured image according to the red square outline.
56
57  Args:
58    img: numpy array; captured image from scene_low_light.
59  Returns:
60    numpy array of the cropped image or the original image if the crop region
61    isn't found.
62  """
63  # To apply k-means clustering, we need to convert the image in to an array
64  # where each row represents a pixel in the image, and each column is a feature
65  # In this case, the feature represents the RGB channels of the pixel
66  data = img.reshape((-1, 3))
67  data = np.float32(data)
68
69  k_means_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
70                      _K_MEANS_ITERATIONS, _K_MEANS_EPSILON)
71  _, labels, centers = cv2.kmeans(data, _NUM_CLUSTERS, None, k_means_criteria,
72                                  _K_MEANS_ITERATIONS,
73                                  cv2.KMEANS_RANDOM_CENTERS)
74  # Find the cluster closest to red
75  min_dist = float('inf')
76  closest_cluster_index = -1
77  for index, center in enumerate(centers):
78    dist = np.linalg.norm(center - np.array(_RED_BGR_COLOR))
79    if dist < min_dist:
80      min_dist = dist
81      closest_cluster_index = index
82
83  target_label = closest_cluster_index
84
85  # create a mask using the data associated with the cluster closest to red
86  mask = labels.flatten() == target_label
87  mask = mask.reshape((img.shape[0], img.shape[1]))
88  mask = mask.astype(np.uint8)
89
90  contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL,
91                                 cv2.CHAIN_APPROX_SIMPLE)
92
93  max_area = 20
94  max_box = None
95
96  # Find the largest box that is closest to square
97  for c in contours:
98    x, y, w, h = cv2.boundingRect(c)
99    aspect_ratio = w / h
100    if _MIN_ASPECT_RATIO < aspect_ratio < _MAX_ASPECT_RATIO:
101      area = w * h
102      if area > max_area:
103        max_area = area
104        max_box = (x, y, w, h)
105
106  # If the box is found then return the cropped image
107  # otherwise the original image is returned
108  if max_box:
109    x, y, w, h = max_box
110    cropped_img = img[
111        y+_CROP_PADDING:y+h-_CROP_PADDING,
112        x+_CROP_PADDING:x+w-_CROP_PADDING
113    ]
114    return cropped_img
115
116  return img
117
118
119def _find_boxes(image):
120  """Finds boxes in the captured image for computing luminance.
121
122  The captured image should be of scene_low_light.png. The boxes are detected
123  by finding the contours by applying a threshold followed erosion.
124
125  Args:
126    image: numpy array; the captured image.
127  Returns:
128    array; an array of boxes, where each box is (x, y, w, h).
129  """
130  gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
131  blur = cv2.GaussianBlur(gray, (3, 3), 0)
132
133  thresh = cv2.adaptiveThreshold(
134      blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 31, -5)
135
136  kernel = np.ones((3, 3), np.uint8)
137  eroded = cv2.erode(thresh, kernel, iterations=1)
138
139  contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL,
140                                 cv2.CHAIN_APPROX_SIMPLE)
141  boxes = []
142
143  # Filter out boxes that are too small or too large
144  # and boxes that are not square
145  img_hw_size_max = max(image.shape[0], image.shape[1])
146  box_min_size = int(round(img_hw_size_max * _BOX_MIN_SIZE_RATIO, 0))
147  if box_min_size == 0:
148    raise AssertionError('Minimum box size calculated was 0. Check cropped '
149                         'image size.')
150  box_max_size = int(img_hw_size_max * _BOX_MAX_SIZE_RATIO)
151  for c in contours:
152    x, y, w, h = cv2.boundingRect(c)
153    aspect_ratio = w / h
154    if (w > box_min_size and h > box_min_size and
155        w < box_max_size and h < box_max_size and
156        _MIN_ASPECT_RATIO < aspect_ratio < _MAX_ASPECT_RATIO):
157      boxes.append((x, y, w, h))
158  return boxes
159
160
161def _correct_image_rotation(img, regions):
162  """Corrects the captured image orientation.
163
164  The captured image should be of scene_low_light.png. The darkest square
165  must appear in the bottom right and the brightest square must appear in
166  the bottom left. This is necessary in order to traverse the hilbert
167  ordered squares to return a darkest to brightest ordering.
168
169  Args:
170    img: numpy array; the original image captured.
171    regions: the tuple of (box, luminance) computed for each square
172      in the image.
173  Returns:
174    numpy array; image in the corrected orientation.
175  """
176  corner_brightness = {
177      _KEY_TOP_LEFT: regions[2][1],
178      _KEY_BOTTOM_LEFT: regions[5][1],
179      _KEY_TOP_RIGHT: regions[14][1],
180      _KEY_BOTTOM_RIGHT: regions[17][1],
181  }
182
183  darkest_corner = ('', float('inf'))
184  brightest_corner = ('', float('-inf'))
185
186  for corner, luminance in corner_brightness.items():
187    if luminance < darkest_corner[1]:
188      darkest_corner = (corner, luminance)
189    if luminance > brightest_corner[1]:
190      brightest_corner = (corner, luminance)
191
192  if darkest_corner == brightest_corner:
193    raise AssertionError('The captured image failed to detect the location '
194                         'of the darkest and brightest squares.')
195
196  if darkest_corner[0] == _KEY_TOP_LEFT:
197    if brightest_corner[0] == _KEY_BOTTOM_LEFT:
198      # rotate 90 CW and then flip vertically
199      img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
200      img = cv2.flip(img, 0)
201    elif brightest_corner[0] == _KEY_TOP_RIGHT:
202      # flip both vertically and horizontally
203      img = cv2.flip(img, -1)
204    else:
205      raise AssertionError('The captured image failed to detect the location '
206                           'of the brightest square.')
207  elif darkest_corner[0] == _KEY_BOTTOM_LEFT:
208    if brightest_corner[0] == _KEY_TOP_LEFT:
209      # rotate 90 CCW
210      img = cv2.rotate(img, cv2.ROTATE_90_COUNTERCLOCKWISE)
211    elif brightest_corner[0] == _KEY_BOTTOM_RIGHT:
212      # flip horizontally
213      img = cv2.flip(img, 1)
214    else:
215      raise AssertionError('The captured image failed to detect the location '
216                           'of the brightest square.')
217  elif darkest_corner[0] == _KEY_TOP_RIGHT:
218    if brightest_corner[0] == _KEY_TOP_LEFT:
219      # flip vertically
220      img = cv2.flip(img, 0)
221    elif brightest_corner[0] == _KEY_BOTTOM_RIGHT:
222      # rotate 90 CW
223      img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
224    else:
225      raise AssertionError('The captured image failed to detect the location '
226                           'of the brightest square.')
227  elif darkest_corner[0] == _KEY_BOTTOM_RIGHT:
228    if brightest_corner[0] == _KEY_BOTTOM_LEFT:
229      # correct orientation
230      pass
231    elif brightest_corner[0] == _KEY_TOP_RIGHT:
232      # rotate 90 and flip horizontally
233      img = cv2.rotate(img, cv2.ROTATE_90_CLOCKWISE)
234      img = cv2.flip(img, 1)
235    else:
236      raise AssertionError('The captured image failed to detect the location '
237                           'of the brightest square.')
238  return img
239
240
241def _compute_luminance_regions(image, boxes):
242  """Compute the luminance for each box in scene_low_light.
243
244  Args:
245    image: numpy array; captured image.
246    boxes: array; array of boxes where each box is (x, y, w, h).
247  Returns:
248    Array of tuples where each tuple is (box, luminance).
249  """
250  intensities = []
251  for b in boxes:
252    x, y, w, h = b
253    padding = min(w, h) * _BOX_PADDING_RATIO
254    left = int(x + padding)
255    top = int(y + padding)
256    right = int(x + w - padding)
257    bottom = int(y + h - padding)
258    box = image[top:bottom, left:right]
259    box_xyz = cv2.cvtColor(box, cv2.COLOR_BGR2XYZ)
260    intensity = int(np.mean(box_xyz[1]))
261    intensities.append((b, intensity))
262  return intensities
263
264
265def _draw_luminance(image, intensities):
266  """Draws the luminance for each box in scene_low_light. Useful for debugging.
267
268  Args:
269    image: numpy array; captured image.
270    intensities: array; array of tuples (box, luminance intensity).
271  """
272  for (b, intensity) in intensities:
273    x, y, w, h = b
274    padding = min(w, h) * _BOX_PADDING_RATIO
275    left = int(x + padding)
276    top = int(y + padding)
277    right = int(x + w - padding)
278    bottom = int(y + h - padding)
279    cv2.rectangle(image, (left, top), (right, bottom), _BOUNDING_BOX_COLOR, 2)
280    cv2.putText(image, f'{intensity}', (x, y - 10),
281                cv2.FONT_HERSHEY_PLAIN, 1, _TEXT_COLOR, 1, 2)
282
283
284def _compute_avg(results):
285  """Computes the average luminance of the first 6 boxes.
286
287  The boxes are part of scene_low_light.
288
289  Args:
290    results: A list of tuples where each tuple is (box, luminance).
291  Returns:
292    float; The average luminance of the first 6 boxes.
293  """
294  luminance_values = [luminance for _, luminance in results[:6]]
295  avg = sum(luminance_values) / len(luminance_values)
296  return avg
297
298
299def _compute_avg_delta_of_successive_boxes(results):
300  """Computes the delta of successive boxes & takes the average of the first 5.
301
302  The boxes are part of scene_low_light.
303
304  Args:
305    results: A list of tuples where each tuple is (box, luminance).
306  Returns:
307    float; The average of the first 5 deltas of successive boxes.
308  """
309  luminance_values = [luminance for _, luminance in results[:6]]
310  delta = [luminance_values[i] - luminance_values[i - 1]
311           for i in range(1, len(luminance_values))]
312  avg = sum(delta) / len(delta)
313  return avg
314
315
316def _plot_results(results, file_stem):
317  """Plots the computed luminance for each box in scene_low_light.
318
319  Args:
320    results: A list of tuples where each tuple is (box, luminance).
321    file_stem: The output file where the plot is saved.
322  """
323  luminance_values = [luminance for _, luminance in results]
324  box_labels = [f'Box {i + 1}' for i in range(len(results))]
325
326  plt.figure(figsize=(10, 6))
327  plt.plot(box_labels, luminance_values, marker='o', linestyle='-', color='b')
328  plt.scatter(box_labels, luminance_values, color='r')
329
330  plt.title('Luminance for each Box')
331  plt.xlabel('Boxes')
332  plt.ylabel('Luminance (pixel intensity)')
333  plt.grid('True')
334  plt.xticks(rotation=45)
335  plt.savefig(f'{file_stem}_luminance_plot.png', dpi=300)
336  plt.close()
337
338
339def _plot_successive_difference(results, file_stem):
340  """Plots the successive difference in luminance between each box.
341
342  The boxes are part of scene_low_light.
343
344  Args:
345    results: A list of tuples where each tuple is (box, luminance).
346    file_stem: The output file where the plot is saved.
347  """
348  luminance_values = [luminance for _, luminance in results]
349  delta = [luminance_values[i] - luminance_values[i - 1]
350           for i in range(1, len(luminance_values))]
351  box_labels = [f'Box {i} to Box {i + 1}' for i in range(1, len(results))]
352
353  plt.figure(figsize=(10, 6))
354  plt.plot(box_labels, delta, marker='o', linestyle='-', color='b')
355  plt.scatter(box_labels, delta, color='r')
356
357  plt.title('Difference in Luminance Between Successive Boxes')
358  plt.xlabel('Box Transition')
359  plt.ylabel('Luminance Difference')
360  plt.grid('True')
361  plt.xticks(rotation=45)
362  file = f'{file_stem}_luminance_difference_between_successive_boxes_plot.png'
363  plt.savefig(file, dpi=300)
364  plt.close()
365
366
367def _sort_by_columns(regions):
368  """Sort the regions by columns and then by row within each column.
369
370  These regions are part of scene_low_light.
371
372  Args:
373    regions: The tuple of (box, luminance) of each square.
374  Returns:
375    array; an array of tuples of (box, luminance) sorted by columns then by row
376      within each column.
377  """
378  # The input is 20 elements. The first two and last two elements represent the
379  # 4 boxes on the outside used for diagnostics. Boxes in indices 2 through 17
380  # represent the elements in the 4x4 grid.
381
382  # Sort all elements by column
383  col_sorted = sorted(regions, key=lambda r: r[0][0])
384
385  # Sort elements within each column by row
386  result = []
387  result.extend(sorted(col_sorted[:2], key=lambda r: r[0][1]))
388
389  for i in range(4):
390    # take 4 rows per column and then sort the rows
391    # skip the first two elements
392    offset = i*4+2
393    col = col_sorted[offset:(offset+4)]
394    result.extend(sorted(col, key=lambda r: r[0][1]))
395
396  result.extend(sorted(col_sorted[-2:], key=lambda r: r[0][1]))
397  return result
398
399
400def analyze_low_light_scene_capture(
401    file_stem,
402    img,
403    avg_luminance_threshold=_LOW_LIGHT_BOOST_AVG_LUMINANCE_THRESH,
404    avg_delta_luminance_threshold=_LOW_LIGHT_BOOST_AVG_DELTA_LUMINANCE_THRESH):
405  """Analyze a captured frame to check if it meets low light scene criteria.
406
407  The capture is cropped first, then detects for boxes, and then computes the
408  luminance of each box.
409
410  Args:
411    file_stem: The file prefix for results saved.
412    img: numpy array; The captured image loaded by cv2 as and available for
413      analysis.
414    avg_luminance_threshold: minimum average luminance of the first 6 boxes.
415    avg_delta_luminance_threshold: minimum average difference in luminance
416      of the first 5 successive boxes of luminance.
417  """
418  cv2.imwrite(f'{file_stem}_original.jpg', img)
419  img = _crop(img)
420  cv2.imwrite(f'{file_stem}_cropped.jpg', img)
421  boxes = _find_boxes(img)
422  if len(boxes) != _EXPECTED_NUM_OF_BOXES:
423    raise AssertionError('The captured image failed to detect the expected '
424                         'number of boxes. '
425                         'Check the captured image to see if the image was '
426                         'correctly captured and try again. '
427                         f'Actual: {len(boxes)}, '
428                         f'Expected: {_EXPECTED_NUM_OF_BOXES}')
429
430  regions = _compute_luminance_regions(img, boxes)
431
432  # Sorted so each column is read left to right
433  sorted_regions = _sort_by_columns(regions)
434  img = _correct_image_rotation(img, sorted_regions)
435  cv2.imwrite(f'{file_stem}_rotated.jpg', img)
436
437  # The orientation of the image may have changed which will affect the
438  # coordinates of the squares. Therefore, locate the squares, recompute the
439  # regions, and sort again
440  boxes = _find_boxes(img)
441  regions = _compute_luminance_regions(img, boxes)
442  sorted_regions = _sort_by_columns(regions)
443
444  _draw_luminance(img, regions)
445  cv2.imwrite(f'{file_stem}_result.jpg', img)
446
447  # Reorder this so the regions are increasing in luminance according to the
448  # Hilbert curve arrangement pattern of the grid
449  # See scene_low_light_reference.png which indicates the order of each
450  # box
451  hilbert_ordered = [
452      sorted_regions[17],
453      sorted_regions[13],
454      sorted_regions[12],
455      sorted_regions[16],
456      sorted_regions[15],
457      sorted_regions[14],
458      sorted_regions[10],
459      sorted_regions[11],
460      sorted_regions[7],
461      sorted_regions[6],
462      sorted_regions[2],
463      sorted_regions[3],
464      sorted_regions[4],
465      sorted_regions[8],
466      sorted_regions[9],
467      sorted_regions[5],
468  ]
469  _plot_results(hilbert_ordered, file_stem)
470  _plot_successive_difference(hilbert_ordered, file_stem)
471  avg = _compute_avg(hilbert_ordered)
472  delta_avg = _compute_avg_delta_of_successive_boxes(hilbert_ordered)
473  test_name = os.path.basename(file_stem)
474
475  # the following print statements are necessary for telemetry
476  # do not convert to logging.debug
477  print(f'{test_name}_avg_luma: {avg:.2f}')
478  print(f'{test_name}_delta_avg_luma: {delta_avg:.2f}')
479  chart_luma_values = [v[1] for v in hilbert_ordered]
480  print(f'{test_name}_chart_luma: {chart_luma_values}')
481
482  logging.debug('average luminance of the 6 boxes: %.2f', avg)
483  logging.debug('average difference in luminance of 5 successive boxes: %.2f',
484                delta_avg)
485  if avg < avg_luminance_threshold:
486    raise AssertionError('Average luminance of the first 6 boxes did not '
487                         'meet minimum requirements for low light scene '
488                         'criteria. '
489                         f'Actual: {avg:.2f}, '
490                         f'Expected: {avg_luminance_threshold}')
491  if delta_avg < avg_delta_luminance_threshold:
492    raise AssertionError('The average difference in luminance of the first 5 '
493                         'successive boxes did not meet minimum requirements '
494                         'for low light scene criteria. '
495                         f'Actual: {delta_avg:.2f}, '
496                         f'Expected: {avg_delta_luminance_threshold}')
497