CameraITS: Parameterizes sensor fusion tests

Allows specification of the output directory for test files, the number
of times to run the test, and the size of the captured images.

Test: test_sensor_fusion.py and run_sensor_fusion_box.py
Bug: 62838524
Change-Id: I612308ad38ffaaf7a763500626b149ac0f185046
diff --git a/apps/CameraITS/tests/sensor_fusion/test_sensor_fusion.py b/apps/CameraITS/tests/sensor_fusion/test_sensor_fusion.py
index a8b9863..f3a7ba4 100644
--- a/apps/CameraITS/tests/sensor_fusion/test_sensor_fusion.py
+++ b/apps/CameraITS/tests/sensor_fusion/test_sensor_fusion.py
@@ -34,10 +34,10 @@
 
 # Capture 210 VGA frames (which is 7s at 30fps)
 N = 210
-W,H = 640,480
-FEATURE_MARGIN = H * 0.20 / 2 # Only take feature points from the center 20%
-                              # so that the rotation measured have much less
-                              # of rolling shutter effect
+W, H = 640, 480
+FEATURE_MARGIN = 0.20  # Only take feature points from the center 20%
+                       # so that the rotation measured have much less of rolling
+                       # shutter effect
 
 MIN_FEATURE_PTS = 30          # Minimum number of feature points required to
                               # perform rotation analysis
@@ -82,22 +82,31 @@
     The instructions for running this test are in the SensorFusion.pdf file in
     the same directory as this test.
 
-    The command-line argument "replay" may be optionally provided. Without this
-    argument, the test will collect a new set of camera+gyro data from the
-    device and then analyze it (and it will also dump this data to files in the
-    current directory). If the "replay" argument is provided, then the script
-    will instead load the dumped data from a previous run and analyze that
-    instead. This can be helpful for developers who are digging for additional
-    information on their measurements.
+    Command line arguments:
+        replay:   Without this argument, the test will collect a new set of
+                  camera+gyro data from the device and then analyze it (and it
+                  will also dump this data to files in the current directory).  If
+                  the "replay" argument is provided, then the script will instead
+                  load the dumped data from a previous run and analyze that
+                  instead. This can be helpful for developers who are digging for
+                  additional information on their measurements.
+        img_size: Comma-separated dimensions of captured images (defaults to
+                  640x480). Ex: 'img_size=<width>,<height>'
     """
 
+    w, h = W, H
+    for s in sys.argv[1:]:
+        if s[:9] == "img_size=" and len(s) > 9:
+            # Split by comma and convert each dimension to int.
+            [w, h] = map(int, s[9:].split(','))
+
     # Collect or load the camera+gyro data. All gyro events as well as camera
     # timestamps are in the "events" dictionary, and "frames" is a list of
     # RGB images as numpy arrays.
     if "replay" not in sys.argv:
-        events, frames = collect_data()
+        events, frames = collect_data(w, h)
     else:
-        events, frames = load_data()
+        events, frames, _, h = load_data()
 
     # Sanity check camera timestamps are enclosed by sensor timestamps
     # This will catch bugs where camera and gyro timestamps go completely out
@@ -124,7 +133,7 @@
 
     # Compute the camera rotation displacements (rad) between each pair of
     # adjacent frames.
-    cam_rots = get_cam_rotations(frames, events["facing"])
+    cam_rots = get_cam_rotations(frames, events["facing"], h)
     if max(abs(cam_rots)) < THRESH_MIN_ROT:
         print "Device wasn't moved enough"
         assert(0)
@@ -270,7 +279,7 @@
     gyro_rots = numpy.array(gyro_rots)
     return gyro_rots
 
-def get_cam_rotations(frames, facing):
+def get_cam_rotations(frames, facing, h):
     """Get the rotations of the camera between each pair of frames.
 
     Takes N frames and returns N-1 angular displacements corresponding to the
@@ -278,6 +287,8 @@
 
     Args:
         frames: List of N images (as RGB numpy arrays).
+        facing: Direction camera is facing
+        h:      Pixel height of each frame
 
     Returns:
         Array of N-1 camera rotation measurements (rad).
@@ -287,8 +298,9 @@
         frame = (frame * 255.0).astype(numpy.uint8)
         gframes.append(cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY))
     rots = []
-    ymin = H/2 - FEATURE_MARGIN
-    ymax = H/2 + FEATURE_MARGIN
+
+    ymin = h*(1-FEATURE_MARGIN)/2
+    ymax = h*(1+FEATURE_MARGIN)/2
     for i in range(1,len(gframes)):
         gframe0 = gframes[i-1]
         gframe1 = gframes[i]
@@ -346,6 +358,8 @@
     Returns:
         events: Dictionary containing all gyro events and cam timestamps.
         frames: List of RGB images as numpy arrays.
+        w:      Pixel width of frames
+        h:      Pixel height of frames
     """
     with open("%s_events.txt"%(NAME), "r") as f:
         events = json.loads(f.read())
@@ -355,14 +369,18 @@
         img = Image.open("%s_frame%03d.png"%(NAME,i))
         w,h = img.size[0:2]
         frames.append(numpy.array(img).reshape(h,w,3) / 255.0)
-    return events, frames
+    return events, frames, w, h
 
-def collect_data():
+def collect_data(w, h):
     """Capture a new set of data from the device.
 
     Captures both motion data and camera frames, while the user is moving
     the device in a proscribed manner.
 
+    Args:
+        w:   Pixel width of frames
+        h:   Pixel height of frames
+
     Returns:
         events: Dictionary containing all gyro events and cam timestamps.
         frames: List of RGB images as numpy arrays.
@@ -387,13 +405,13 @@
             print "Unknown lens facing", facing
             assert(0)
 
-        fmt = {"format":"yuv", "width":W, "height":H}
+        fmt = {"format":"yuv", "width":w, "height":h}
         s,e,_,_,_ = cam.do_3a(get_results=True, do_af=False)
         req = its.objects.manual_capture_request(s, e)
         fps = 30
         req["android.control.aeTargetFpsRange"] = [fps, fps]
         print "Capturing %dx%d with sens. %d, exp. time %.1fms" % (
-                W, H, s, e*NSEC_TO_MSEC)
+                w, h, s, e*NSEC_TO_MSEC)
         caps = cam.do_capture([req]*N, fmt)
 
         # Get the gyro events.
diff --git a/apps/CameraITS/tools/run_sensor_fusion_box.py b/apps/CameraITS/tools/run_sensor_fusion_box.py
index 6eb8ceb..1722690 100644
--- a/apps/CameraITS/tools/run_sensor_fusion_box.py
+++ b/apps/CameraITS/tools/run_sensor_fusion_box.py
@@ -20,9 +20,7 @@
 import time
 
 import its.device
-# from its.device import ItsSession
 
-NUM_RUNS = 2
 SCENE_NAME = 'sensor_fusion'
 SKIP_RET_CODE = 101
 TEST_NAME = 'test_sensor_fusion'
@@ -36,27 +34,40 @@
 
     Script should be run from the top-level CameraITS directory.
 
-    Command line Arguments:
-        camera:  the camera(s) to be tested. Use comma to separate multiple
-                 camera Ids. Ex: 'camera=0,1' or 'camera=1'
-        device:  the device id for adb
-        rotator: string for rotator id in for vid:pid:ch
+    Command line arguments:
+        camera:   camera(s) to be tested. Use comma to separate multiple
+                  camera Ids. Ex: 'camera=0,1' or 'camera=1'
+        device:   device id for adb
+        num_runs: number of times to repeat the test
+        rotator:  string for rotator id in for vid:pid:ch
+        tmp_dir:  location of temp directory for output files
+        img_size: Comma-separated dimensions of captured images (defaults to
+                  640x480). Ex: 'img_size=<width>,<height>'
     """
 
     camera_id = '0'
+    num_runs = 1
     rotator_ids = 'default'
+    tmp_dir = None
+    img_size = '640,480'
     for s in sys.argv[1:]:
         if s[:7] == 'camera=' and len(s) > 7:
             camera_id = s[7:]
+        elif s[:9] == 'num_runs=' and len(s) > 9:
+            num_runs = int(s[9:])
         elif s[:8] == 'rotator=' and len(s) > 8:
             rotator_ids = s[8:]
+        elif s[:8] == 'tmp_dir=' and len(s) > 8:
+            tmp_dir = s[8:]
+        elif s[:9] == 'img_size=' and len(s) > 9:
+            img_size = s[9:]
 
     if camera_id not in ['0', '1']:
         print 'Need to specify camera 0 or 1'
         sys.exit()
 
     # Make output directories to hold the generated files.
-    tmpdir = tempfile.mkdtemp()
+    tmpdir = tempfile.mkdtemp(dir=tmp_dir)
     print 'Saving output files to:', tmpdir, '\n'
 
     device_id = its.device.get_device_id()
@@ -68,18 +79,21 @@
         rotator_id_arg = 'rotator=' + rotator_ids
     print 'Preparing to run sensor_fusion on camera', camera_id
 
+    img_size_arg = 'size=' + img_size
+    print 'Image dimensions are ' + img_size
+
     os.mkdir(os.path.join(tmpdir, camera_id))
 
-    # Run test multiple times, capturing stdout and stderr.
+    # Run test "num_runs" times, capturing stdout and stderr.
     numpass = 0
     numfail = 0
     numskip = 0
-    for i in range(NUM_RUNS):
+    for i in range(num_runs):
         os.mkdir(os.path.join(tmpdir, camera_id, SCENE_NAME+'_'+str(i)))
         cmd = ('python tools/rotation_rig.py rotator=%s' % rotator_ids)
         subprocess.Popen(cmd.split())
         cmd = ['python', os.path.join(TEST_DIR, TEST_NAME+'.py'),
-               device_id_arg, camera_id_arg, rotator_id_arg]
+               device_id_arg, camera_id_arg, rotator_id_arg, img_size_arg]
         outdir = os.path.join(tmpdir, camera_id, SCENE_NAME+'_'+str(i))
         outpath = os.path.join(outdir, TEST_NAME+'_stdout.txt')
         errpath = os.path.join(outdir, TEST_NAME+'_stderr.txt')
@@ -102,7 +116,7 @@
         print msg
 
     test_result = '%d / %d tests passed (%.1f%%)' % (
-        numpass+numskip, NUM_RUNS, 100.0*float(numpass+numskip)/NUM_RUNS)
+        numpass+numskip, num_runs, 100*float(numpass+numskip)/num_runs)
     print test_result
 
 if __name__ == '__main__':