Added capacity for multiple rectangle labels; bug fixes

1 files changed, 61 insertions, 9 deletions

diff --git a/makeCrop.py b/makeCrop.py
index 8865dad..19ce369 100755
--- a/makeCrop.py
+++ b/makeCrop.py
@@ -1,22 +1,21 @@
 #!/usr/bin/env python3
-# PYTHON_ARGCOMPLETE_OK
 
-import argparse, argcomplete
+import argparse
 
-parser = argparse.ArgumentParser(description='Crop and scale a video based on bounding boxes')
+parser = argparse.ArgumentParser(description='Crop and scale a video based on bounding boxes', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 parser.add_argument('boxes', help='Path to csv file where bounding boxes are saved')
 parser.add_argument('video', help='Path to video to crop')
 parser.add_argument('--intermediary', help='Path to intermediary video if size differs')
+parser.add_argument('--no-interpolate', help='Do not interpolate over gaps between box updates', action='store_true')
 parser.add_argument('--label', help='Boxes label (uses all if not provided)')
-parser.add_argument('--save-stem', help='Location to save stem.npz and stem.avi. Default: out', default='out')
+parser.add_argument('--save-stem', help='Location to save stem.npz and stem.avi', default='out')
 parser.add_argument('--skip-zero-boxes', help='Skip boxes with zero size', action='store_true')
+parser.add_argument('--outputWH', help='Output width and height', nargs=2, default=[64,64], type=int)
 parser.add_argument('--debug', help='Save a debugging video as well', action='store_true')
 
-argcomplete.autocomplete(parser)
 args = parser.parse_args()
 
-OUTW = 64
-OUTH = 64
+OUTW, OUTH = args.outputWH
 
 import csv
 with open(args.boxes) as f:
@@ -47,6 +46,51 @@ boxes = [[b[0], float(b[1]), int(int(b[2]) * zoomX), int(int(b[3]) * zoomY), int
 if args.skip_zero_boxes:
     boxes = [b for b in boxes if b[2] != b[4] and b[3] != b[5]]
 
+def alignDims(box: list, whRatio: float = 1.0) -> list:
+    '''Make a bounding box dimentions adhere to given ratio
+
+    Arguments:
+      box: The bounding box, formatted [x1, y1, x2, y2]
+      whRatio: The desired ratio of width to height
+
+    Returns:
+      list: The adjusted box
+    '''
+    coords1, coords2 = [0, 2], [1, 3]
+    if (box[coords2[1]] - box[coords2[0]]) * whRatio > (box[coords1[1]] - box[coords1[0]]):
+        coords1, coords2 = coords2, coords1
+        whRatio = 1 / whRatio
+    # coords1 > whRatio * coords2, but we want them equal
+    d = int((box[coords1[1]] - box[coords1[0]]) / whRatio - (box[coords2[1]] - box[coords2[0]]))
+    box[coords2[0]] -= int(d/2)
+    box[coords2[1]] += int(d/2) + (d % 2)
+    return box
+
+def shift(box: list, w: int, h: int) -> list:
+    '''Shift a bounding box to be within w,h bounds
+
+    Arguments:
+      box: The bounding box, formatted [x1, y1, x2, y2]
+      w: The width of the frame (i.e., maximum x value)
+      h: The height of the frame (i.e., maximum y value)
+
+    Returns:
+      list: The shifted box
+    '''
+    bounds = [[0, int(w)], [0, int(h)]]
+    coords = [[0, 2], [1, 3]]
+    for b, c in zip(bounds, coords):
+        s = 0
+        if box[c[0]] < b[0]:
+            s = -box[c[0]]
+        elif box[c[1]] > b[1]:
+            s = b[1] - box[c[1]]
+        box[c[0]] += s
+        box[c[1]] += s
+    return box
+
+boxes = [box[0:2] + shift(alignDims(box[2:], OUTW/OUTH), width, height) for box in boxes]
+
 import numpy as np
 outArry = []
 outVid = cv2.VideoWriter(args.save_stem + '.avi', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (OUTW, OUTH))
@@ -67,8 +111,14 @@ while cap.isOpened():
     # Make boxes[0] the most recent timestamp without going over current time
     while len(boxes) >= 2 and boxes[1][1] <= time:
         boxes = boxes[1:]
-    x1, x2 = sorted([boxes[0][2], boxes[0][4]])
-    y1, y2 = sorted([boxes[0][3], boxes[0][5]])
+    if len(boxes) == 1: # Stop when we run out
+        break
+    box = boxes[0]
+    if not args.no_interpolate:
+        weight = (boxes[1][1] - time) / (boxes[1][1] - boxes[0][1])
+        box = [boxes[0][0], time] + [int(boxes[0][i] * weight + boxes[1][i] * (1-weight)) for i in range(2, 6)]
+    x1, x2 = sorted([box[2], box[4]])
+    y1, y2 = sorted([box[3], box[5]])
     cropped = frame[y1:y2, x1:x2]
     if cropped.size < 1:
         resized = np.zeros((OUTH, OUTW, 3), dtype=cropped.dtype)
@@ -88,5 +138,7 @@ while cap.isOpened():
 
 bar.finish()
 
+print(f'Ending at time = {time}')
+
 # Save out array
 np.savez_compressed(args.save_stem + '.npz', np.array(outArry))