1 files changed, 92 insertions, 0 deletions
diff --git a/makeCrop.py b/makeCrop.py
new file mode 100755
index 0000000..8865dad
--- /dev/null
+++ b/makeCrop.py
@@ -0,0 +1,92 @@
+#!/usr/bin/env python3
+# PYTHON_ARGCOMPLETE_OK
+
+import argparse, argcomplete
+
+parser = argparse.ArgumentParser(description='Crop and scale a video based on bounding boxes')
+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('--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('--skip-zero-boxes', help='Skip boxes with zero size', action='store_true')
+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
+
+import csv
+with open(args.boxes) as f:
+    boxes = list(csv.reader(f))
+
+if args.label:
+    boxes = [b for b in boxes if b[0] == args.label]
+
+import cv2
+cap = cv2.VideoCapture(args.video)
+width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
+height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
+fps = cap.get(cv2.CAP_PROP_FPS)
+
+zoomX = 1.0
+zoomY = 1.0
+if args.intermediary:
+    inter = cv2.VideoCapture(args.intermediary)
+    interWidth = inter.get(cv2.CAP_PROP_FRAME_WIDTH)
+    interHeight = inter.get(cv2.CAP_PROP_FRAME_HEIGHT)
+    zoomX = width / interWidth
+    zoomY = height / interHeight
+    inter.release()
+
+# Each box is [str(label), float(time), int(x1), int(y1), int(x2), int(y2)]
+boxes = [[b[0], float(b[1]), int(int(b[2]) * zoomX), int(int(b[3]) * zoomY), int(int(b[4]) * zoomX), int(int(b[5]) * zoomY)] for b in boxes]
+
+if args.skip_zero_boxes:
+    boxes = [b for b in boxes if b[2] != b[4] and b[3] != b[5]]
+
+import numpy as np
+outArry = []
+outVid = cv2.VideoWriter(args.save_stem + '.avi', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (OUTW, OUTH))
+
+if args.debug:
+    outDebug = cv2.VideoWriter(args.save_stem + '-debug.avi', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (int(OUTW * height / OUTH + width), int(height)))
+
+# This may take a bit, so do a progress bar
+from progress.bar import IncrementalBar
+bar = IncrementalBar('Frames Processed', max=cap.get(cv2.CAP_PROP_FRAME_COUNT), suffix='%(index)d/%(max)d - %(eta)d s')
+
+frameNum = 0
+while cap.isOpened():
+    ret, frame = cap.read()
+    if not ret:
+        break
+    time = frameNum / fps
+    # 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]])
+    cropped = frame[y1:y2, x1:x2]
+    if cropped.size < 1:
+        resized = np.zeros((OUTH, OUTW, 3), dtype=cropped.dtype)
+    else:
+        resized = cv2.resize(cropped, (OUTH, OUTW), interpolation=cv2.INTER_CUBIC)
+    outArry.append(resized)
+    outVid.write(resized)
+    if args.debug:
+        # Draw box on frame
+        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        # Resize "resized"
+        scaledUp = cv2.resize(resized, (int(OUTW * height / OUTH), int(height)), interpolation=cv2.INTER_AREA)
+        concatted = cv2.hconcat([frame, scaledUp])
+        outDebug.write(concatted)
+    bar.next()
+    frameNum += 1
+
+bar.finish()
+
+# Save out array
+np.savez_compressed(args.save_stem + '.npz', np.array(outArry))