1# Copyright 2020 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 sensor_fusion hardware rig."""
15
16
17import bisect
18import codecs
19import logging
20import math
21import os
22import struct
23import time
24
25import cv2
26from matplotlib import pylab
27import matplotlib.pyplot
28import numpy as np
29import scipy.spatial
30import serial
31from serial.tools import list_ports
32
33import camera_properties_utils
34import image_processing_utils
35
36# Constants for Rotation Rig
37ARDUINO_ANGLE_MAX = 180.0  # degrees
38ARDUINO_ANGLES_SENSOR_FUSION = (0, 90)  # degrees
39ARDUINO_ANGLES_STABILIZATION = (10, 25)  # degrees
40ARDUINO_BAUDRATE = 9600
41ARDUINO_CMD_LENGTH = 3
42ARDUINO_CMD_TIME = 2.0 * ARDUINO_CMD_LENGTH / ARDUINO_BAUDRATE  # round trip
43ARDUINO_MOVE_TIME_SENSOR_FUSION = 2  # seconds
44ARDUINO_MOVE_TIME_STABILIZATION = 0.3  # seconds
45ARDUINO_PID = 0x0043
46ARDUINO_SERVO_SPEED_MAX = 255
47ARDUINO_SERVO_SPEED_MIN = 1
48ARDUINO_SERVO_SPEED_SENSOR_FUSION = 20
49ARDUINO_SERVO_SPEED_STABILIZATION = 10
50ARDUINO_SERVO_SPEED_STABILIZATION_TABLET = 20
51ARDUINO_SPEED_START_BYTE = 253
52ARDUINO_START_BYTE = 255
53ARDUINO_START_NUM_TRYS = 5
54ARDUINO_START_TIMEOUT = 300  # seconds
55ARDUINO_STRING = 'Arduino'
56ARDUINO_TEST_CMD = (b'\x01', b'\x02', b'\x03')
57ARDUINO_VALID_CH = ('1', '2', '3', '4', '5', '6')
58ARDUINO_VIDS = (0x2341, 0x2a03)
59
60CANAKIT_BAUDRATE = 115200
61CANAKIT_CMD_TIME = 0.05  # seconds (found experimentally)
62CANAKIT_DATA_DELIMITER = '\r\n'
63CANAKIT_PID = 0xfc73
64CANAKIT_SEND_TIMEOUT = 0.02  # seconds
65CANAKIT_SET_CMD = 'REL'
66CANAKIT_SLEEP_TIME = 2  # seconds (for full 90 degree rotation)
67CANAKIT_VALID_CMD = ('ON', 'OFF')
68CANAKIT_VALID_CH = ('1', '2', '3', '4')
69CANAKIT_VID = 0x04d8
70
71HS755HB_ANGLE_MAX = 202.0  # throw for rotation motor in degrees
72
73# From test_sensor_fusion
74_FEATURE_MARGIN = 0.20  # Only take feature points from center 20% so that
75                        # rotation measured has less rolling shutter effect.
76_FEATURE_PTS_MIN = 30  # Min number of feature pts to perform rotation analysis.
77# cv2.goodFeatures to track.
78# 'POSTMASK' is the measurement method in all previous versions of Android.
79# 'POSTMASK' finds best features on entire frame and then masks the features
80# to the vertical center FEATURE_MARGIN for the measurement.
81# 'PREMASK' is a new measurement that is used when FEATURE_PTS_MIN is not
82# found in frame. This finds the best 2*FEATURE_PTS_MIN in the FEATURE_MARGIN
83# part of the frame.
84_CV2_FEATURE_PARAMS_POSTMASK = dict(maxCorners=240,
85                                    qualityLevel=0.3,
86                                    minDistance=7,
87                                    blockSize=7)
88_CV2_FEATURE_PARAMS_PREMASK = dict(maxCorners=2*_FEATURE_PTS_MIN,
89                                   qualityLevel=0.3,
90                                   minDistance=7,
91                                   blockSize=7)
92_GYRO_SAMP_RATE_MIN = 100.0  # Samples/second: min gyro sample rate.
93_CV2_LK_PARAMS = dict(winSize=(15, 15),
94                      maxLevel=2,
95                      criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT,
96                                10, 0.03))  # cv2.calcOpticalFlowPyrLK params.
97_ROTATION_PER_FRAME_MIN = 0.001  # rads/s
98_GYRO_ROTATION_PER_SEC_MAX = 2.0  # rads/s
99_R_SQUARED_TOLERANCE = 0.01  # tolerance for polynomial fitting r^2
100_SHIFT_DOMAIN_RADIUS = 5  # limited domain centered around best shift
101
102# unittest constants
103_COARSE_FIT_RANGE = 20  # Range area around coarse fit to do optimization.
104_CORR_TIME_OFFSET_MAX = 50  # ms max shift to try and match camera/gyro times.
105_CORR_TIME_OFFSET_STEP = 0.5  # ms step for shifts.
106
107# Unit translators
108_MSEC_TO_NSEC = 1000000
109_NSEC_TO_SEC = 1E-9
110_SEC_TO_NSEC = int(1/_NSEC_TO_SEC)
111_RADS_TO_DEGS = 180/math.pi
112
113_NUM_GYRO_PTS_TO_AVG = 20
114
115
116def polynomial_from_coefficients(coefficients):
117  """Return a polynomial function from a coefficient list, highest power first.
118
119  Args:
120    coefficients: list of coefficients (float)
121  Returns:
122    Function in the form of a*x^n + b*x^(n - 1) + ... + constant
123  """
124  def polynomial(x):
125    n = len(coefficients)
126    return sum(coefficients[i] * x ** (n - i - 1) for i in range(n))
127  return polynomial
128
129
130def smallest_absolute_minimum_of_polynomial(coefficients):
131  """Return the smallest minimum by absolute value from a coefficient list.
132
133  Args:
134    coefficients: list of coefficients (float)
135  Returns:
136    Smallest local minimum (by absolute value) on the function (float)
137  """
138  first_derivative = np.polyder(coefficients, m=1)
139  second_derivative = np.polyder(coefficients, m=2)
140  extrema = np.roots(first_derivative)
141  smallest_absolute_minimum = None
142  for extremum in extrema:
143    if np.polyval(second_derivative, extremum) > 0:
144      if smallest_absolute_minimum is None or abs(extremum) < abs(
145          smallest_absolute_minimum):
146        smallest_absolute_minimum = extremum
147  if smallest_absolute_minimum is None:
148    raise AssertionError(
149        f'No minima were found on function described by {coefficients}.')
150  return smallest_absolute_minimum
151
152
153def serial_port_def(name):
154  """Determine the serial port and open.
155
156  Args:
157    name: string of device to locate (ie. 'Arduino', 'Canakit' or 'Default')
158  Returns:
159    serial port object
160  """
161  serial_port = None
162  devices = list_ports.comports()
163  for device in devices:
164    if not (device.vid and device.pid):  # Not all comm ports have vid and pid
165      continue
166    if name.lower() == 'arduino':
167      if (device.vid in ARDUINO_VIDS and device.pid == ARDUINO_PID):
168        logging.debug('Arduino: %s', str(device))
169        serial_port = device.device
170        return serial.Serial(serial_port, ARDUINO_BAUDRATE, timeout=1)
171
172    elif name.lower() in ('canakit', 'default'):
173      if (device.vid == CANAKIT_VID and device.pid == CANAKIT_PID):
174        logging.debug('Canakit: %s', str(device))
175        serial_port = device.device
176        return serial.Serial(serial_port, CANAKIT_BAUDRATE,
177                             timeout=CANAKIT_SEND_TIMEOUT,
178                             parity=serial.PARITY_EVEN,
179                             stopbits=serial.STOPBITS_ONE,
180                             bytesize=serial.EIGHTBITS)
181  raise ValueError(f'{name} device not connected.')
182
183
184def canakit_cmd_send(canakit_serial_port, cmd_str):
185  """Wrapper for sending serial command to Canakit.
186
187  Args:
188    canakit_serial_port: port to write for canakit
189    cmd_str: str; value to send to device.
190  """
191  try:
192    logging.debug('writing port...')
193    canakit_serial_port.write(CANAKIT_DATA_DELIMITER.encode())
194    time.sleep(CANAKIT_CMD_TIME)  # This is critical for relay.
195    canakit_serial_port.write(cmd_str.encode())
196
197  except IOError as io_error:
198    raise IOError(
199        f'Port {CANAKIT_VID}:{CANAKIT_PID} is not open!') from io_error
200
201
202def canakit_set_relay_channel_state(canakit_port, ch, state):
203  """Set Canakit relay channel and state.
204
205  Waits CANAKIT_SLEEP_TIME for rotation to occur.
206
207  Args:
208    canakit_port: serial port object for the Canakit port.
209    ch: string for channel number of relay to set. '1', '2', '3', or '4'
210    state: string of either 'ON' or 'OFF'
211  """
212  logging.debug('Setting relay state %s', state)
213  if ch in CANAKIT_VALID_CH and state in CANAKIT_VALID_CMD:
214    canakit_cmd_send(canakit_port, CANAKIT_SET_CMD + ch + '.' + state + '\r\n')
215    time.sleep(CANAKIT_SLEEP_TIME)
216  else:
217    logging.debug('Invalid ch (%s) or state (%s), no command sent.', ch, state)
218
219
220def arduino_read_cmd(port):
221  """Read back Arduino command from serial port."""
222  cmd = []
223  for _ in range(ARDUINO_CMD_LENGTH):
224    cmd.append(port.read())
225  return cmd
226
227
228def arduino_send_cmd(port, cmd):
229  """Send command to serial port."""
230  for i in range(ARDUINO_CMD_LENGTH):
231    port.write(cmd[i])
232
233
234def arduino_loopback_cmd(port, cmd):
235  """Send command to serial port."""
236  arduino_send_cmd(port, cmd)
237  time.sleep(ARDUINO_CMD_TIME)
238  return arduino_read_cmd(port)
239
240
241def establish_serial_comm(port):
242  """Establish connection with serial port."""
243  logging.debug('Establishing communication with %s', port.name)
244  trys = 1
245  hex_test = convert_to_hex(ARDUINO_TEST_CMD)
246  logging.debug(' test tx: %s %s %s', hex_test[0], hex_test[1], hex_test[2])
247  start = time.time()
248  while time.time() < start + ARDUINO_START_TIMEOUT:
249    try:
250      cmd_read = arduino_loopback_cmd(port, ARDUINO_TEST_CMD)
251    except serial.serialutil.SerialException as _:
252      logging.debug('Port in use, trying again...')
253      continue
254    hex_read = convert_to_hex(cmd_read)
255    logging.debug(' test rx: %s %s %s', hex_read[0], hex_read[1], hex_read[2])
256    if cmd_read != list(ARDUINO_TEST_CMD):
257      trys += 1
258    else:
259      logging.debug(' Arduino comm established after %d try(s)', trys)
260      break
261  else:
262    raise AssertionError(f'Arduino comm not established after {trys} tries '
263                         f'and {ARDUINO_START_TIMEOUT} seconds')
264
265
266def convert_to_hex(cmd):
267  return [('%0.2x' % int(codecs.encode(x, 'hex_codec'), 16) if x else '--')
268          for x in cmd]
269
270
271def arduino_rotate_servo_to_angle(ch, angle, serial_port, move_time):
272  """Rotate servo to the specified angle.
273
274  Args:
275    ch: str; servo to rotate in ARDUINO_VALID_CH
276    angle: int; servo angle to move to
277    serial_port: object; serial port
278    move_time: int; time in seconds
279  """
280  if angle < 0 or angle > ARDUINO_ANGLE_MAX:
281    logging.debug('Angle must be between 0 and %d.', ARDUINO_ANGLE_MAX)
282    angle = 0
283    if angle > ARDUINO_ANGLE_MAX:
284      angle = ARDUINO_ANGLE_MAX
285
286  cmd = [struct.pack('B', i) for i in [ARDUINO_START_BYTE, int(ch), angle]]
287  arduino_send_cmd(serial_port, cmd)
288  time.sleep(move_time)
289
290
291def arduino_rotate_servo(ch, angles, move_time, serial_port):
292  """Rotate servo through 'angles'.
293
294  Args:
295    ch: str; servo to rotate
296    angles: list of ints; servo angles to move to
297    move_time: int; time required to allow for arduino movement
298    serial_port: object; serial port
299  """
300
301  for angle in angles:
302    angle_norm = int(round(angle*ARDUINO_ANGLE_MAX/HS755HB_ANGLE_MAX, 0))
303    arduino_rotate_servo_to_angle(ch, angle_norm, serial_port, move_time)
304
305
306def rotation_rig(rotate_cntl, rotate_ch, num_rotations, angles, servo_speed,
307                 move_time, arduino_serial_port):
308  """Rotate the phone n times using rotate_cntl and rotate_ch defined.
309
310  rotate_ch is hard wired and must be determined from physical setup.
311  If using Arduino, serial port must be initialized and communication must be
312  established before rotation.
313
314  Args:
315    rotate_cntl: str to identify 'arduino', 'canakit' or 'external' controller.
316    rotate_ch: str to identify rotation channel number.
317    num_rotations: int number of rotations.
318    angles: list of ints; servo angle to move to.
319    servo_speed: int number of move speed between [1, 255].
320    move_time: int time required to allow for arduino movement.
321    arduino_serial_port: optional initialized serial port object
322  """
323
324  logging.debug('Controller: %s, ch: %s', rotate_cntl, rotate_ch)
325  if arduino_serial_port:
326    # initialize servo at origin
327    logging.debug('Moving servo to origin')
328    arduino_rotate_servo_to_angle(rotate_ch, 0, arduino_serial_port, 1)
329
330    # set servo speed
331    set_servo_speed(rotate_ch, servo_speed, arduino_serial_port, delay=0)
332  elif rotate_cntl.lower() == 'canakit':
333    canakit_serial_port = serial_port_def('Canakit')
334  elif rotate_cntl.lower() == 'external':
335    logging.info('External rotation control.')
336  else:
337    logging.info('No rotation rig defined. Manual test: rotate phone by hand.')
338
339  # rotate phone
340  logging.debug('Rotating phone %dx', num_rotations)
341  for _ in range(num_rotations):
342    if rotate_cntl == 'arduino':
343      arduino_rotate_servo(rotate_ch, angles, move_time, arduino_serial_port)
344    elif rotate_cntl == 'canakit':
345      canakit_set_relay_channel_state(canakit_serial_port, rotate_ch, 'ON')
346      canakit_set_relay_channel_state(canakit_serial_port, rotate_ch, 'OFF')
347  logging.debug('Finished rotations')
348  if rotate_cntl == 'arduino':
349    logging.debug('Moving servo to origin')
350    arduino_rotate_servo_to_angle(rotate_ch, 0, arduino_serial_port, 1)
351
352
353def set_servo_speed(ch, servo_speed, serial_port, delay=0):
354  """Set servo to specified speed.
355
356  Args:
357    ch: str; servo to turn on in ARDUINO_VALID_CH
358    servo_speed: int; value of speed between 1 and 255
359    serial_port: object; serial port
360    delay: int; time in seconds
361  """
362  logging.debug('Servo speed: %d', servo_speed)
363  if servo_speed < ARDUINO_SERVO_SPEED_MIN:
364    logging.debug('Servo speed must be >= %d.', ARDUINO_SERVO_SPEED_MIN)
365    servo_speed = ARDUINO_SERVO_SPEED_MIN
366  elif servo_speed > ARDUINO_SERVO_SPEED_MAX:
367    logging.debug('Servo speed must be <= %d.', ARDUINO_SERVO_SPEED_MAX)
368    servo_speed = ARDUINO_SERVO_SPEED_MAX
369
370  cmd = [struct.pack('B', i) for i in [ARDUINO_SPEED_START_BYTE,
371                                       int(ch), servo_speed]]
372  arduino_send_cmd(serial_port, cmd)
373  time.sleep(delay)
374
375
376def calc_max_rotation_angle(rotations, sensor_type):
377  """Calculates the max angle of deflection from rotations.
378
379  Args:
380    rotations: numpy array of rotation per event
381    sensor_type: string 'Camera' or 'Gyro'
382
383  Returns:
384    maximum angle of rotation for the given rotations
385  """
386  rotations *= _RADS_TO_DEGS
387  rotations_sum = np.cumsum(rotations)
388  rotation_max = max(rotations_sum)
389  rotation_min = min(rotations_sum)
390  logging.debug('%s min: %.2f, max %.2f rotation (degrees)',
391                sensor_type, rotation_min, rotation_max)
392  logging.debug('%s max rotation: %.2f degrees',
393                sensor_type, (rotation_max-rotation_min))
394  return rotation_max-rotation_min
395
396
397def get_gyro_rotations(gyro_events, cam_times):
398  """Get the rotation values of the gyro.
399
400  Integrates the gyro data between each camera frame to compute an angular
401  displacement.
402
403  Args:
404    gyro_events: List of gyro event objects.
405    cam_times: Array of N camera times, one for each frame.
406
407  Returns:
408    Array of N-1 gyro rotation measurements (rads/s).
409  """
410  gyro_times = np.array([e['time'] for e in gyro_events])
411  all_gyro_rots = np.array([e['z'] for e in gyro_events])
412  gyro_rots = []
413  if gyro_times[0] > cam_times[0] or gyro_times[-1] < cam_times[-1]:
414    raise AssertionError('Gyro times do not bound camera times! '
415                         f'gyro: {gyro_times[0]:.0f} -> {gyro_times[-1]:.0f} '
416                         f'cam: {cam_times[0]} -> {cam_times[-1]} (ns).')
417
418  # Integrate the gyro data between each pair of camera frame times.
419  for i_cam in range(len(cam_times)-1):
420    # Get the window of gyro samples within the current pair of frames.
421    # Note: bisect always picks first gyro index after camera time.
422    t_cam0 = cam_times[i_cam]
423    t_cam1 = cam_times[i_cam+1]
424    i_gyro_window0 = bisect.bisect(gyro_times, t_cam0)
425    i_gyro_window1 = bisect.bisect(gyro_times, t_cam1)
426    gyro_sum = 0
427
428    # Integrate samples within the window.
429    for i_gyro in range(i_gyro_window0, i_gyro_window1):
430      gyro_val = all_gyro_rots[i_gyro+1]
431      t_gyro0 = gyro_times[i_gyro]
432      t_gyro1 = gyro_times[i_gyro+1]
433      t_gyro_delta = (t_gyro1 - t_gyro0) * _NSEC_TO_SEC
434      gyro_sum += gyro_val * t_gyro_delta
435
436    # Handle the fractional intervals at the sides of the window.
437    for side, i_gyro in enumerate([i_gyro_window0-1, i_gyro_window1]):
438      gyro_val = all_gyro_rots[i_gyro+1]
439      t_gyro0 = gyro_times[i_gyro]
440      t_gyro1 = gyro_times[i_gyro+1]
441      t_gyro_delta = (t_gyro1 - t_gyro0) * _NSEC_TO_SEC
442      if side == 0:
443        f = (t_cam0 - t_gyro0) / (t_gyro1 - t_gyro0)
444        frac_correction = gyro_val * t_gyro_delta * (1.0 - f)
445        gyro_sum += frac_correction
446      else:
447        f = (t_cam1 - t_gyro0) / (t_gyro1 - t_gyro0)
448        frac_correction = gyro_val * t_gyro_delta * f
449        gyro_sum += frac_correction
450    gyro_rots.append(gyro_sum)
451  gyro_rots = np.array(gyro_rots)
452  return gyro_rots
453
454
455def procrustes_rotation(x, y):
456  """Performs a Procrustes analysis to conform points in x to y.
457
458  Procrustes analysis determines a linear transformation (translation,
459  reflection, orthogonal rotation and scaling) of the points in y to best
460  conform them to the points in matrix x, using the sum of squared errors
461  as the metric for fit criterion.
462
463  Args:
464    x: Target coordinate matrix
465    y: Input coordinate matrix
466
467  Returns:
468    The rotation component of the transformation that maps x to y.
469  """
470  x0 = (x-x.mean(0)) / np.sqrt(((x-x.mean(0))**2.0).sum())
471  y0 = (y-y.mean(0)) / np.sqrt(((y-y.mean(0))**2.0).sum())
472  u, _, vt = np.linalg.svd(np.dot(x0.T, y0), full_matrices=False)
473  return np.dot(vt.T, u.T)
474
475
476def get_cam_rotations(frames, facing, h, file_name_stem,
477                      start_frame, stabilized_video=False):
478  """Get the rotations of the camera between each pair of frames.
479
480  Takes N frames and returns N-1 angular displacements corresponding to the
481  rotations between adjacent pairs of frames, in radians.
482  Only takes feature points from center so that rotation measured has less
483  rolling shutter effect.
484  Requires FEATURE_PTS_MIN to have enough data points for accurate measurements.
485  Uses FEATURE_PARAMS for cv2 to identify features in checkerboard images.
486  Ensures camera rotates enough if not calling with stabilized video.
487
488  Args:
489    frames: List of N images (as RGB numpy arrays).
490    facing: Direction camera is facing.
491    h: Pixel height of each frame.
492    file_name_stem: file name stem including location for data.
493    start_frame: int; index to start at
494    stabilized_video: Boolean; if called with stabilized video
495
496  Returns:
497    numpy array of N-1 camera rotation measurements (rad).
498  """
499  gframes = []
500  for frame in frames:
501    frame = (frame * 255.0).astype(np.uint8)  # cv2 uses [0, 255]
502    gframes.append(cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY))
503  num_frames = len(gframes)
504  logging.debug('num_frames: %d', num_frames)
505  # create mask
506  ymin = int(h * (1 - _FEATURE_MARGIN) / 2)
507  ymax = int(h * (1 + _FEATURE_MARGIN) / 2)
508  pre_mask = np.zeros_like(gframes[0])
509  pre_mask[ymin:ymax, :] = 255
510
511  for masking in ['post', 'pre']:  # Do post-masking (original) method 1st
512    logging.debug('Using %s masking method', masking)
513    rotations = []
514    for i in range(1, num_frames):
515      j = i - 1
516      gframe0 = gframes[j]
517      gframe1 = gframes[i]
518      if masking == 'post':
519        p0 = cv2.goodFeaturesToTrack(
520            gframe0, mask=None, **_CV2_FEATURE_PARAMS_POSTMASK)
521        post_mask = (p0[:, 0, 1] >= ymin) & (p0[:, 0, 1] <= ymax)
522        p0_filtered = p0[post_mask]
523      else:
524        p0_filtered = cv2.goodFeaturesToTrack(
525            gframe0, mask=pre_mask, **_CV2_FEATURE_PARAMS_PREMASK)
526      num_features = len(p0_filtered)
527      if num_features < _FEATURE_PTS_MIN:
528        for pt in np.rint(p0_filtered).astype(int):
529          x, y = pt[0][0], pt[0][1]
530          cv2.circle(frames[j], (x, y), 3, (100, 255, 255), -1)
531        image_processing_utils.write_image(
532            frames[j], f'{file_name_stem}_features{j+start_frame:03d}.png')
533        msg = (f'Not enough features in frame {j+start_frame}. Need at least '
534               f'{_FEATURE_PTS_MIN} features, got {num_features}.')
535        if masking == 'pre':
536          raise AssertionError(msg)
537        else:
538          logging.debug(msg)
539          break
540      else:
541        logging.debug('Number of features in frame %s is %d',
542                      str(j+start_frame).zfill(3), num_features)
543      p1, st, _ = cv2.calcOpticalFlowPyrLK(gframe0, gframe1, p0_filtered, None,
544                                           **_CV2_LK_PARAMS)
545      tform = procrustes_rotation(p0_filtered[st == 1], p1[st == 1])
546      if facing == camera_properties_utils.LENS_FACING['BACK']:
547        rotation = -math.atan2(tform[0, 1], tform[0, 0])
548      elif facing == camera_properties_utils.LENS_FACING['FRONT']:
549        rotation = math.atan2(tform[0, 1], tform[0, 0])
550      else:
551        raise AssertionError(f'Unknown lens facing: {facing}.')
552      rotations.append(rotation)
553      if i == 1:
554        # Save debug visualization of features that are being
555        # tracked in the first frame.
556        frame = frames[j]
557        for x, y in np.rint(p0_filtered[st == 1]).astype(int):
558          cv2.circle(frame, (x, y), 3, (100, 255, 255), -1)
559        image_processing_utils.write_image(
560            frame, f'{file_name_stem}_features{j+start_frame:03d}.png')
561    if i == num_frames-1:
562      logging.debug('Correct num of frames found: %d', i)
563      break  # exit if enough features in all frames
564  if i != num_frames-1:
565    raise AssertionError('Neither method found enough features in all frames')
566
567  rotations = np.array(rotations)
568  rot_per_frame_max = max(abs(rotations))
569  logging.debug('Max rotation in frame: %.2f degrees',
570                rot_per_frame_max*_RADS_TO_DEGS)
571  if stabilized_video:
572    logging.debug('Skipped camera rotation check due to stabilized video.')
573  else:
574    if rot_per_frame_max < _ROTATION_PER_FRAME_MIN:
575      raise AssertionError(f'Device not moved enough: {rot_per_frame_max:.3f} '
576                           f'movement. THRESH: {_ROTATION_PER_FRAME_MIN} rads.')
577    else:
578      logging.debug('Device movement exceeds %.2f degrees',
579                    _ROTATION_PER_FRAME_MIN*_RADS_TO_DEGS)
580  return rotations
581
582
583def get_best_alignment_offset(cam_times, cam_rots, gyro_events, degree=2):
584  """Find the best offset to align the camera and gyro motion traces.
585
586  This function integrates the shifted gyro data between camera samples
587  for a range of candidate shift values, and returns the shift that
588  result in the best correlation.
589
590  Uses a correlation distance metric between the curves, where a smaller
591  value means that the curves are better-correlated.
592
593  Fits a curve to the correlation distance data to measure the minima more
594  accurately, by looking at the correlation distances within a range of
595  +/- 10ms from the measured best score; note that this will use fewer
596  than the full +/- 10 range for the curve fit if the measured score
597  (which is used as the center of the fit) is within 10ms of the edge of
598  the +/- 50ms candidate range.
599
600  Args:
601    cam_times: Array of N camera times, one for each frame.
602    cam_rots: Array of N-1 camera rotation displacements (rad).
603    gyro_events: List of gyro event objects.
604    degree: Degree of polynomial
605
606  Returns:
607    Best alignment offset(ms), fit coefficients, candidates, and distances.
608  """
609  # Measure the correlation distance over defined shift
610  shift_candidates = np.arange(-_CORR_TIME_OFFSET_MAX,
611                               _CORR_TIME_OFFSET_MAX+_CORR_TIME_OFFSET_STEP,
612                               _CORR_TIME_OFFSET_STEP).tolist()
613  spatial_distances = []
614  for shift in shift_candidates:
615    shifted_cam_times = cam_times + shift*_MSEC_TO_NSEC
616    gyro_rots = get_gyro_rotations(gyro_events, shifted_cam_times)
617    spatial_distance = scipy.spatial.distance.correlation(cam_rots, gyro_rots)
618    logging.debug('shift %.1fms spatial distance: %.5f', shift,
619                  spatial_distance)
620    spatial_distances.append(spatial_distance)
621
622  best_corr_dist = min(spatial_distances)
623  coarse_best_shift = shift_candidates[spatial_distances.index(best_corr_dist)]
624  logging.debug('Best shift without fitting is %.4f ms', coarse_best_shift)
625
626  # Fit a polynomial around coarse_best_shift to extract best fit
627  i = spatial_distances.index(best_corr_dist)
628  i_poly_fit_min = i - _COARSE_FIT_RANGE
629  i_poly_fit_max = i + _COARSE_FIT_RANGE + 1
630  shift_candidates = shift_candidates[i_poly_fit_min:i_poly_fit_max]
631  spatial_distances = spatial_distances[i_poly_fit_min:i_poly_fit_max]
632  logging.debug('Polynomial degree: %d', degree)
633  fit_coeffs, residuals, _, _, _ = np.polyfit(
634      shift_candidates, spatial_distances, degree, full=True
635  )
636  logging.debug('Fit coefficients: %s', fit_coeffs)
637  logging.debug('Residuals: %s', residuals)
638  total_sum_of_squares = np.sum(
639      (spatial_distances - np.mean(spatial_distances)) ** 2
640  )
641  # Calculate r-squared on the entire domain for debugging
642  r_squared = 1 - residuals[0] / total_sum_of_squares
643  logging.debug('r^2 on the entire domain: %f', r_squared)
644
645  # Calculate r-squared near the best shift
646  domain_around_best_shift = [coarse_best_shift - _SHIFT_DOMAIN_RADIUS,
647                              coarse_best_shift + _SHIFT_DOMAIN_RADIUS]
648  logging.debug('Calculating r^2 on the limited domain of [%f, %f]',
649                domain_around_best_shift[0], domain_around_best_shift[1])
650  small_shifts_and_distances = [
651      (x, y)
652      for x, y in zip(shift_candidates, spatial_distances)
653      if domain_around_best_shift[0] <= x <= domain_around_best_shift[1]
654  ]
655  small_shift_candidates, small_spatial_distances = zip(
656      *small_shifts_and_distances
657  )
658  logging.debug('Shift candidates on limited domain: %s',
659                small_shift_candidates)
660  logging.debug('Spatial distances on limited domain: %s',
661                small_spatial_distances)
662  limited_residuals = np.sum(
663      (np.polyval(fit_coeffs, small_shift_candidates) - small_spatial_distances)
664      ** 2
665  )
666  logging.debug('Residuals on limited domain: %s', limited_residuals)
667  limited_total_sum_of_squares = np.sum(
668      (small_spatial_distances - np.mean(small_spatial_distances)) ** 2
669  )
670  limited_r_squared = 1 - limited_residuals / limited_total_sum_of_squares
671  logging.debug('r^2 on limited domain: %f', limited_r_squared)
672
673  # Calculate exact_best_shift (x where y is minimum of parabola)
674  exact_best_shift = smallest_absolute_minimum_of_polynomial(fit_coeffs)
675
676  if abs(coarse_best_shift - exact_best_shift) > 2.0:
677    raise AssertionError(
678        f'Test failed. Bad fit to time-shift curve. Coarse best shift: '
679        f'{coarse_best_shift}, Exact best shift: {exact_best_shift}.')
680
681  # Check fit of polynomial near the best shift
682  if not math.isclose(limited_r_squared, 1, abs_tol=_R_SQUARED_TOLERANCE):
683    logging.debug('r-squared on domain [%f, %f] was %f, expected 1.0, '
684                  'ATOL: %f',
685                  domain_around_best_shift[0], domain_around_best_shift[1],
686                  limited_r_squared, _R_SQUARED_TOLERANCE)
687    return None
688
689  return exact_best_shift, fit_coeffs, shift_candidates, spatial_distances
690
691
692def plot_camera_rotations(cam_rots, start_frame, video_quality,
693                          plot_name_stem):
694  """Plot the camera rotations.
695
696  Args:
697   cam_rots: np array of camera rotations angle per frame
698   start_frame: int value of start frame
699   video_quality: str for video quality identifier
700   plot_name_stem: str (with path) of what to call plot
701  """
702
703  pylab.figure(video_quality)
704  frames = range(start_frame, len(cam_rots)+start_frame)
705  pylab.title(f'Camera rotation vs frame {video_quality}')
706  pylab.plot(frames, cam_rots*_RADS_TO_DEGS, '-ro', label='x')
707  pylab.xlabel('frame #')
708  pylab.ylabel('camera rotation (degrees)')
709  matplotlib.pyplot.savefig(f'{plot_name_stem}_cam_rots.png')
710  pylab.close(video_quality)
711
712
713def plot_gyro_events(gyro_events, plot_name, log_path):
714  """Plot x, y, and z on the gyro events.
715
716  Samples are grouped into NUM_GYRO_PTS_TO_AVG groups and averaged to minimize
717  random spikes in data.
718
719  Args:
720    gyro_events: List of gyroscope events.
721    plot_name:  name of plot(s).
722    log_path: location to save data.
723  """
724
725  nevents = (len(gyro_events) // _NUM_GYRO_PTS_TO_AVG) * _NUM_GYRO_PTS_TO_AVG
726  gyro_events = gyro_events[:nevents]
727  times = np.array([(e['time'] - gyro_events[0]['time']) * _NSEC_TO_SEC
728                    for e in gyro_events])
729  x = np.array([e['x'] for e in gyro_events])
730  y = np.array([e['y'] for e in gyro_events])
731  z = np.array([e['z'] for e in gyro_events])
732
733  # Group samples into size-N groups & average each together to minimize random
734  # spikes in data.
735  times = times[_NUM_GYRO_PTS_TO_AVG//2::_NUM_GYRO_PTS_TO_AVG]
736  x = x.reshape(nevents//_NUM_GYRO_PTS_TO_AVG, _NUM_GYRO_PTS_TO_AVG).mean(1)
737  y = y.reshape(nevents//_NUM_GYRO_PTS_TO_AVG, _NUM_GYRO_PTS_TO_AVG).mean(1)
738  z = z.reshape(nevents//_NUM_GYRO_PTS_TO_AVG, _NUM_GYRO_PTS_TO_AVG).mean(1)
739
740  pylab.figure(plot_name)
741  # x & y on same axes
742  pylab.subplot(2, 1, 1)
743  pylab.title(f'{plot_name}(mean of {_NUM_GYRO_PTS_TO_AVG} pts)')
744  pylab.plot(times, x, 'r', label='x')
745  pylab.plot(times, y, 'g', label='y')
746  pylab.ylim([np.amin(z)/4, np.amax(z)/4])  # zoom in 4x from z axis
747  pylab.ylabel('gyro x,y movement (rads/s)')
748  pylab.legend()
749
750  # z on separate axes
751  pylab.subplot(2, 1, 2)
752  pylab.plot(times, z, 'b', label='z')
753  pylab.ylim([np.amin(z), np.amax(z)])
754  pylab.xlabel('time (seconds)')
755  pylab.ylabel('gyro z movement (rads/s)')
756  pylab.legend()
757  file_name = os.path.join(log_path, plot_name)
758  matplotlib.pyplot.savefig(f'{file_name}_gyro_events.png')
759  pylab.close(plot_name)
760
761  z_max = max(abs(z))
762  logging.debug('z_max: %.3f', z_max)
763  if z_max > _GYRO_ROTATION_PER_SEC_MAX:
764    raise AssertionError(
765        f'Phone moved too rapidly! Please confirm controller firmware. '
766        f'Max: {z_max:.3f}, TOL: {_GYRO_ROTATION_PER_SEC_MAX} rads/s')
767
768
769def conv_acceleration_to_movement(gyro_events, video_delay_time):
770  """Convert gyro_events time and speed to movement during video time.
771
772  Args:
773    gyro_events: sorted dict of entries with 'time', 'x', 'y', and 'z'
774    video_delay_time: time at which video starts (and the video's duration)
775
776  Returns:
777    'z' acceleration converted to movement for times around VIDEO playing.
778  """
779  gyro_times = np.array([e['time'] for e in gyro_events])
780  gyro_speed = np.array([e['z'] for e in gyro_events])
781  gyro_time_min = gyro_times[0]
782  logging.debug('gyro start time: %dns', gyro_time_min)
783  logging.debug('gyro stop time: %dns', gyro_times[-1])
784  gyro_rotations = []
785  video_time_start = gyro_time_min + video_delay_time *_SEC_TO_NSEC
786  video_time_stop = video_time_start + video_delay_time *_SEC_TO_NSEC
787  logging.debug('video start time: %dns', video_time_start)
788  logging.debug('video stop time: %dns', video_time_stop)
789
790  for i, t in enumerate(gyro_times):
791    if video_time_start <= t <= video_time_stop:
792      gyro_rotations.append((gyro_times[i]-gyro_times[i-1])/_SEC_TO_NSEC *
793                            gyro_speed[i])
794  return np.array(gyro_rotations)
795