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"""Verify that frames from UW and W cameras are not distorted."""
15
16import collections
17import logging
18import os
19import cv2
20import math
21import numpy as np
22
23from cv2 import aruco
24from mobly import test_runner
25
26import its_base_test
27import camera_properties_utils
28import image_processing_utils
29import its_session_utils
30import preview_processing_utils
31
32_ACCURACY = 0.001
33_ARUCO_COUNT = 8
34_ARUCO_DIST_TOL = 0.15
35_ARUCO_SIZE = (3, 3)
36_ASPECT_RATIO_4_3 = 4/3
37_CH_FULL_SCALE = 255
38_CHESSBOARD_CORNERS = 24
39_CHKR_DIST_TOL = 0.05
40_CROSS_SIZE = 6
41_CROSS_THICKNESS = 1
42_FONT_SCALE = 0.3
43_FONT_THICKNESS = 1
44_GREEN_LIGHT = (80, 255, 80)
45_GREEN_DARK = (0, 190, 0)
46_MAX_ITER = 30
47_NAME = os.path.splitext(os.path.basename(__file__))[0]
48_RED = (255, 0, 0)
49_VALID_CONTROLLERS = ('arduino', 'external')
50_WIDE_ZOOM = 1
51_ZOOM_STEP = 0.5
52_ZOOM_STEP_REDUCTION = 0.1
53_ZOOM_TOL = 0.1
54
55
56def get_chart_coverage(image, corners):
57  """Calculates the chart coverage in the image.
58
59  Args:
60    image: image containing chessboard
61    corners: corners of the chart
62
63  Returns:
64    chart_coverage: percentage of the image covered by chart corners
65    chart_diagonal_pixels: pixel count from the first corner to the last corner
66  """
67  first_corner = corners[0].tolist()[0]
68  logging.debug('first_corner: %s', first_corner)
69  last_corner = corners[-1].tolist()[0]
70  logging.debug('last_corner: %s', last_corner)
71  chart_diagonal_pixels = math.dist(first_corner, last_corner)
72  logging.debug('chart_diagonal_pixels: %s', chart_diagonal_pixels)
73
74  # Calculate chart coverage relative to image diagonal
75  image_diagonal = np.sqrt(image.shape[0]**2 + image.shape[1]**2)
76  logging.debug('image.shape: %s', image.shape)
77  logging.debug('Image diagonal (pixels): %s', image_diagonal)
78  chart_coverage = chart_diagonal_pixels / image_diagonal * 100
79  logging.debug('Chart coverage: %s', chart_coverage)
80
81  return chart_coverage, chart_diagonal_pixels
82
83
84def plot_corners(image, corners, cross_color=_RED, text_color=_RED):
85  """Plot corners to the given image.
86
87  Args:
88    image: image
89    corners: points in the image
90    cross_color: color of cross
91    text_color: color of text
92
93  Returns:
94    image: image with cross and text for each corner
95  """
96  for i, corner in enumerate(corners):
97    x, y = int(corner.ravel()[0]), int(corner.ravel()[1])
98
99    # Draw corner index
100    cv2.putText(image, str(i), (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX,
101                _FONT_SCALE, text_color, _FONT_THICKNESS, cv2.LINE_AA)
102
103  for corner in corners:
104    x, y = corner.ravel()
105
106    # Ensure coordinates are integers and within image boundaries
107    x = max(0, min(int(x), image.shape[1] - 1))
108    y = max(0, min(int(y), image.shape[0] - 1))
109
110    # Draw horizontal line
111    cv2.line(image, (x - _CROSS_SIZE, y), (x + _CROSS_SIZE, y), cross_color,
112             _CROSS_THICKNESS)
113    # Draw vertical line
114    cv2.line(image, (x, y - _CROSS_SIZE), (x, y + _CROSS_SIZE), cross_color,
115             _CROSS_THICKNESS)
116
117  return image
118
119
120def get_ideal_points(pattern_size):
121  """Calculate the ideal points for pattern.
122
123  These are just corners at unit intervals of the same dimensions
124  as pattern_size. Looks like..
125   [[ 0.  0.  0.]
126    [ 1.  0.  0.]
127    [ 2.  0.  0.]
128     ...
129    [21. 23.  0.]
130    [22. 23.  0.]
131    [23. 23.  0.]]
132
133  Args:
134    pattern_size: pattern size. Example (24, 24)
135
136  Returns:
137    ideal_points: corners at unit interval.
138  """
139  ideal_points = np.zeros((pattern_size[0] * pattern_size[1], 3), np.float32)
140  ideal_points[:,:2] = (
141      np.mgrid[0:pattern_size[0], 0:pattern_size[1]].T.reshape(-1, 2)
142  )
143
144  return ideal_points
145
146
147def get_distortion_error(image, corners, ideal_points, rotation_vector,
148                         translation_vector, camera_matrix):
149  """Get distortion error by comparing corners and ideal points.
150
151  compare corners and ideal points to derive the distortion error
152
153  Args:
154    image: image containing chessboard and ArUco
155    corners: corners of the chart. Shape = (number of corners, 1, 2)
156    ideal_points: corners at unit interval. Shape = (number of corners, 3)
157    rotation_vector: rotation vector based on chart's rotation. Shape = (3, 1)
158    translation_vector: translation vector based on chart's rotation.
159                        Shape = (3, 1)
160    camera_matrix: camera intrinsic matrix. Shape = (3, 3)
161
162  Returns:
163    normalized_distortion_error_percentage: normalized distortion error
164      percentage. None if all corners based on pattern_size not found.
165    chart_coverage: percentage of the image covered by corners
166  """
167  chart_coverage, chart_diagonal_pixels = get_chart_coverage(image, corners)
168  logging.debug('Chart coverage: %s', chart_coverage)
169
170  projected_points = cv2.projectPoints(ideal_points, rotation_vector,
171                                       translation_vector, camera_matrix, None)
172  # Reshape projected points to 2D array
173  projected = projected_points[0].reshape(-1, 2)
174  corners_reshaped = corners.reshape(-1, 2)
175  logging.debug('projected: %s', projected)
176
177  plot_corners(image, projected, _GREEN_LIGHT, _GREEN_DARK)
178
179  # Calculate the distortion error
180  distortion_errors = [
181      math.dist(projected_point, corner_point)
182      for projected_point, corner_point in zip(projected, corners_reshaped)
183  ]
184  logging.debug('distortion_error: %s', distortion_errors)
185
186  # Get RMS of error
187  rms_error = math.sqrt(np.mean(np.square(distortion_errors)))
188  logging.debug('RMS distortion error: %s', rms_error)
189
190  # Calculate as a percentage of the chart diagonal
191  normalized_distortion_error_percentage = (
192      rms_error / chart_diagonal_pixels * 100
193  )
194  logging.debug('Normalized percent distortion error: %s',
195                normalized_distortion_error_percentage)
196
197  return normalized_distortion_error_percentage, chart_coverage
198
199
200def get_chessboard_corners(pattern_size, image):
201  """Find chessboard corners from image.
202
203  Args:
204    pattern_size: (int, int) chessboard corners.
205    image: image containing chessboard
206
207  Returns:
208    corners: corners of the chessboard chart
209    ideal_points: ideal pattern of chessboard corners
210                  i.e. points at unit intervals
211  """
212  # Convert the image to grayscale
213  gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
214
215  # Find the checkerboard corners
216  found_corners, corners_pass1 = cv2.findChessboardCorners(gray_image,
217                                                           pattern_size)
218  logging.debug('Found corners: %s', found_corners)
219  logging.debug('corners_pass1: %s', corners_pass1)
220
221  if not found_corners:
222    logging.debug('Chessboard pattern not found.')
223    return None, None
224
225  # Refine corners
226  criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, _MAX_ITER,
227              _ACCURACY)
228  corners = cv2.cornerSubPix(gray_image, corners_pass1, (11, 11), (-1, -1),
229                             criteria)
230  logging.debug('Refined Corners: %s', corners)
231
232  plot_corners(image, corners)
233
234  ideal_points = get_ideal_points(pattern_size)
235  logging.debug('ideal_points: %s', ideal_points)
236
237  return corners, ideal_points
238
239
240def get_aruco_corners(image):
241  """Find ArUco corners from image.
242
243  Args:
244    image: image containing ArUco markers
245
246  Returns:
247    corners: First corner of each ArUco markers in the image.
248             None if expected ArUco corners are not found.
249    ideal_points: ideal pattern of the ArUco marker corners.
250                  None if expected ArUco corners are not found.
251  """
252  # Detect ArUco markers
253  aruco_dict = aruco.getPredefinedDictionary(aruco.DICT_4X4_100)
254  corners, ids, _ = aruco.detectMarkers(image, aruco_dict)
255
256  logging.debug('corners: %s', corners)
257  logging.debug('ids: %s', ids)
258
259  if ids is None:
260    logging.debug('ArUco markers are not found')
261    return None, None
262
263  aruco.drawDetectedMarkers(image, corners, ids, _RED)
264
265  # Convert to numpy array
266  corners = np.concatenate(corners, axis=0).reshape(-1, 4, 2)
267
268  # Extract first corners efficiently
269  corners = corners[:, 0, :]
270  logging.debug('corners: %s', corners)
271
272  # Create marker_dict using efficient vectorization
273  marker_dict = dict(zip(ids.flatten(), corners))
274
275  if len(marker_dict) != _ARUCO_COUNT:
276    logging.debug('%s arUCO markers found instead of %s',
277                  len(ids), _ARUCO_COUNT)
278    return None, None
279
280  # Arrange corners based on ids
281  arranged_corners = np.array([marker_dict[i] for i in range(len(corners))])
282
283  # Add a dimension to match format for cv2.calibrateCamera
284  corners = np.expand_dims(arranged_corners, axis=1)
285  logging.debug('updated corners: %s', corners)
286
287  plot_corners(image, corners)
288
289  ideal_points = get_ideal_points(_ARUCO_SIZE)
290
291  # No ArUco marker in the center, so remove the middle point
292  middle_index = (_ARUCO_SIZE[0] // 2) * _ARUCO_SIZE[1] + (_ARUCO_SIZE[1] // 2)
293  ideal_points = np.delete(ideal_points, middle_index, axis=0)
294  logging.debug('ideal_points: %s', ideal_points)
295
296  return corners, ideal_points
297
298
299def get_preview_frame(dut, cam, preview_size, zoom, log_path):
300  """Captures preview frame at given zoom ratio.
301
302  Args:
303    dut: device under test
304    cam: camera object
305    preview_size: str; preview resolution. ex. '1920x1080'
306    zoom: zoom ratio
307    log_path: str; path for video file directory
308
309  Returns:
310    img_name: the filename of the first captured image
311    capture_result: total capture results of the preview frame
312  """
313  # Define zoom fields such that preview recording is at only one zoom level
314  z_min = zoom
315  z_max = z_min + _ZOOM_STEP - _ZOOM_STEP_REDUCTION
316
317  # Capture preview images over zoom range
318  # TODO: b/343200676 - use do_preview_recording instead of
319  #                     preview_over_zoom_range
320  capture_results, file_list = preview_processing_utils.preview_over_zoom_range(
321      dut, cam, preview_size, z_min, z_max, _ZOOM_STEP, log_path
322  )
323
324  # Get first captured image
325  img_name = file_list[0]
326  capture_result = capture_results[0]
327
328  return img_name, capture_result
329
330
331def add_update_to_filename(file_name, update_str='_update'):
332  """Adds the provided update string to the base name of a file.
333
334  Args:
335    file_name (str): The full path to the file to be modified.
336    update_str (str, optional): The string to insert before the extension
337
338  Returns:
339    file_name: The full path to the new file with the update string added.
340  """
341
342  directory, file_with_ext = os.path.split(file_name)
343  base_name, ext = os.path.splitext(file_with_ext)
344
345  new_file_name = os.path.join(directory, f'{base_name}_{update_str}{ext}')
346
347  return new_file_name
348
349
350def get_distortion_errors(props, img_name):
351  """Calculates the distortion error using checkerboard and ArUco markers.
352
353  Args:
354    props: camera properties object.
355    img_name: image name including complete file path
356
357  Returns:
358    chkr_chart_coverage: normalized distortion error percentage for chessboard
359      corners. None if all corners based on pattern_size not found.
360    chkr_chart_coverage: percentage of the image covered by chessboard chart
361    arc_distortion_error: normalized distortion error percentage for ArUco
362      corners. None if all corners based on pattern_size not found.
363    arc_chart_coverage: percentage of the image covered by ArUco corners
364
365  """
366  image = cv2.imread(img_name)
367  if (props['android.lens.facing'] ==
368      camera_properties_utils.LENS_FACING['FRONT']):
369    image = preview_processing_utils.mirror_preview_image_by_sensor_orientation(
370        props['android.sensor.orientation'], image)
371
372  pattern_size = (_CHESSBOARD_CORNERS, _CHESSBOARD_CORNERS)
373
374  chess_corners, chess_ideal_points = get_chessboard_corners(pattern_size,
375                                                             image)
376  aruco_corners, aruco_ideal_points = get_aruco_corners(image)
377
378  if chess_corners is None:
379    return None, None, None, None
380
381  ideal_points = [chess_ideal_points]
382  image_corners = [chess_corners]
383
384  if aruco_corners is not None:
385    ideal_points.append(aruco_ideal_points)
386    image_corners.append(aruco_corners)
387
388  # Calculate the distortion error
389  # Do this by:
390  # 1) Calibrate the camera from the detected checkerboard points
391  # 2) Project the ideal points, using the camera calibration data.
392  # 3) Except, do not use distortion coefficients so we model ideal pinhole
393  # 4) Calculate the error of the detected corners relative to the ideal
394  # 5) Normalize the average error by the size of the chart
395  calib_flags = (
396      cv2.CALIB_FIX_K1
397      + cv2.CALIB_FIX_K2
398      + cv2.CALIB_FIX_K3
399      + cv2.CALIB_FIX_K4
400      + cv2.CALIB_FIX_K5
401      + cv2.CALIB_FIX_K6
402      + cv2.CALIB_ZERO_TANGENT_DIST
403  )
404  ret, camera_matrix, dist_coeffs, rotation_vectors, translation_vectors = (
405      cv2.calibrateCamera(ideal_points, image_corners, image.shape[:2],
406                          None, None, flags=calib_flags)
407  )
408  logging.debug('Projection error: %s dist_coeffs: %s', ret, dist_coeffs)
409  logging.debug('rotation_vector: %s', rotation_vectors)
410  logging.debug('translation_vector: %s', translation_vectors)
411  logging.debug('matrix: %s', camera_matrix)
412
413  chkr_distortion_error, chkr_chart_coverage = (
414      get_distortion_error(image, chess_corners, chess_ideal_points,
415                           rotation_vectors[0], translation_vectors[0],
416                           camera_matrix)
417  )
418
419  if aruco_corners is not None:
420    arc_distortion_error, arc_chart_coverage = get_distortion_error(
421        image, aruco_corners, aruco_ideal_points, rotation_vectors[1],
422        translation_vectors[1], camera_matrix
423    )
424  else:
425    arc_distortion_error, arc_chart_coverage = None, None
426
427  img_name_update = add_update_to_filename(img_name)
428  image_processing_utils.write_image(image / _CH_FULL_SCALE, img_name_update)
429
430  return (chkr_distortion_error, chkr_chart_coverage,
431          arc_distortion_error, arc_chart_coverage)
432
433
434class PreviewDistortionTest(its_base_test.ItsBaseTest):
435  """Test that frames from UW and W cameras are not distorted.
436
437  Captures preview frames at different zoom levels. If whole chart is visible
438  in the frame, detect the distortion error. Pass the test if distortion error
439  is within the pre-determined TOL.
440  """
441
442  def test_preview_distortion(self):
443    rot_rig = {}
444    log_path = self.log_path
445
446    with its_session_utils.ItsSession(
447        device_id=self.dut.serial,
448        camera_id=self.camera_id,
449        hidden_physical_id=self.hidden_physical_id) as cam:
450
451      props = cam.get_camera_properties()
452      props = cam.override_with_hidden_physical_camera_props(props)
453      camera_properties_utils.skip_unless(
454          camera_properties_utils.zoom_ratio_range(props))
455
456      # Raise error if not FRONT or REAR facing camera
457      camera_properties_utils.check_front_or_rear_camera(props)
458
459      # Initialize rotation rig
460      rot_rig['cntl'] = self.rotator_cntl
461      rot_rig['ch'] = self.rotator_ch
462      if rot_rig['cntl'].lower() not in _VALID_CONTROLLERS:
463        raise AssertionError(
464            f'You must use the {_VALID_CONTROLLERS} controller for {_NAME}.')
465
466      # Determine preview size
467      preview_size = preview_processing_utils.get_max_preview_test_size(
468          cam, self.camera_id, _ASPECT_RATIO_4_3)
469      logging.debug('preview_size: %s', preview_size)
470
471      # Determine test zoom range
472      z_range = props['android.control.zoomRatioRange']
473      logging.debug('z_range: %s', z_range)
474
475      # Collect preview frames and associated capture results
476      PreviewFrameData = collections.namedtuple(
477          'PreviewFrameData', ['img_name', 'capture_result', 'z_level']
478      )
479      preview_frames = []
480      z_levels = [z_range[0]]  # Min zoom
481      if z_range[0] < _WIDE_ZOOM:
482        z_levels.append(_WIDE_ZOOM)
483
484      for z in z_levels:
485        img_name, capture_result = get_preview_frame(
486            self.dut, cam, preview_size, z, log_path
487        )
488        if img_name:
489          frame_data = PreviewFrameData(img_name, capture_result, z)
490          preview_frames.append(frame_data)
491
492      failure_msg = []
493      # Determine distortion error and chart coverage for each frames
494      for frame in preview_frames:
495        img_full_name = f'{os.path.join(log_path, frame.img_name)}'
496        (chkr_distortion_err, chkr_chart_coverage, arc_distortion_err,
497         arc_chart_coverage) = get_distortion_errors(props, img_full_name)
498
499        zoom = float(frame.capture_result['android.control.zoomRatio'])
500        if camera_properties_utils.logical_multi_camera(props):
501          cam_id = frame.capture_result[
502              'android.logicalMultiCamera.activePhysicalId'
503          ]
504        else:
505          cam_id = None
506        logging.debug('Zoom: %.2f, cam_id: %s, img_name: %s',
507                      zoom, cam_id, img_name)
508
509        if math.isclose(zoom, z_levels[0], rel_tol=_ZOOM_TOL):
510          z_str = 'min'
511        else:
512          z_str = 'max'
513
514        # Don't change print to logging. Used for KPI.
515        print(f'{_NAME}_{z_str}_zoom: ', zoom)
516        print(f'{_NAME}_{z_str}_physical_id: ', cam_id)
517        print(f'{_NAME}_{z_str}_chkr_distortion_error: ', chkr_distortion_err)
518        print(f'{_NAME}_{z_str}_chkr_chart_coverage: ', chkr_chart_coverage)
519        print(f'{_NAME}_{z_str}_aruco_distortion_error: ', arc_distortion_err)
520        print(f'{_NAME}_{z_str}_aruco_chart_coverage: ', arc_chart_coverage)
521        logging.debug('%s_%s_zoom: %s', _NAME, z_str, zoom)
522        logging.debug('%s_%s_physical_id: %s', _NAME, z_str, cam_id)
523        logging.debug('%s_%s_chkr_distortion_error: %s', _NAME, z_str,
524                      chkr_distortion_err)
525        logging.debug('%s_%s_chkr_chart_coverage: %s', _NAME, z_str,
526                      chkr_chart_coverage)
527        logging.debug('%s_%s_aruco_distortion_error: %s', _NAME, z_str,
528                      arc_distortion_err)
529        logging.debug('%s_%s_aruco_chart_coverage: %s', _NAME, z_str,
530                      arc_chart_coverage)
531
532        if arc_distortion_err is None:
533          if zoom < _WIDE_ZOOM:
534            failure_msg.append('Unable to find all ArUco markers in '
535                               f'{img_name}')
536            logging.debug(failure_msg[-1])
537        else:
538          if arc_distortion_err > _ARUCO_DIST_TOL:
539            failure_msg.append('ArUco Distortion error '
540                               f'{arc_distortion_err:.3f} is greater than '
541                               f'tolerance {_ARUCO_DIST_TOL}')
542            logging.debug(failure_msg[-1])
543
544        if chkr_distortion_err is None:
545          # Checkerboard corners shall be detected at minimum zoom level
546          failure_msg.append(f'Unable to find full checker board in {img_name}')
547          logging.debug(failure_msg[-1])
548        else:
549          if chkr_distortion_err > _CHKR_DIST_TOL:
550            failure_msg.append('Chess Distortion error '
551                               f'{chkr_distortion_err:.3f} is greater than '
552                               f'tolerance {_CHKR_DIST_TOL}')
553            logging.debug(failure_msg[-1])
554
555      if failure_msg:
556        raise AssertionError(f'{its_session_utils.NOT_YET_MANDATED_MESSAGE}'
557                             f'\n\n{failure_msg}')
558
559if __name__ == '__main__':
560  test_runner.main()
561
562