Switched crop target from npz to hdf5HEADmaster

1 files changed, 17 insertions, 2 deletions

diff --git a/makeCrop.py b/makeCrop.py
index 19ce369..70b2fc2 100755
--- a/makeCrop.py
+++ b/makeCrop.py
@@ -8,22 +8,27 @@ 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')
+parser.add_argument('--save-stem', help='Location to save stem.hdf5 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')
+parser.add_argument('--extraMetadata', help='Path to json-formatted extra metadata (in addition to embedded in video file)')
 
 args = parser.parse_args()
 
 OUTW, OUTH = args.outputWH
 
 import csv
+import json
+import subprocess
 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]
 
+metadata=json.loads(subprocess.check_output(['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_format', '-show_streams', args.video]))
+
 import cv2
 cap = cv2.VideoCapture(args.video)
 width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
@@ -32,7 +37,9 @@ fps = cap.get(cv2.CAP_PROP_FPS)
 
 zoomX = 1.0
 zoomY = 1.0
+
 if args.intermediary:
+    metadata=json.loads(subprocess.check_output(['ffprobe', '-v', 'quiet', '-print_format', 'json', '-show_format', '-show_streams', args.intermediary]))
     inter = cv2.VideoCapture(args.intermediary)
     interWidth = inter.get(cv2.CAP_PROP_FRAME_WIDTH)
     interHeight = inter.get(cv2.CAP_PROP_FRAME_HEIGHT)
@@ -40,6 +47,10 @@ if args.intermediary:
     zoomY = height / interHeight
     inter.release()
 
+if args.extraMetadata:
+    with open(args.extraMetadata) as f:
+        metadata['extra'] = json.load(f)
+
 # 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]
 
@@ -92,6 +103,7 @@ def shift(box: list, w: int, h: int) -> list:
 boxes = [box[0:2] + shift(alignDims(box[2:], OUTW/OUTH), width, height) for box in boxes]
 
 import numpy as np
+import h5py
 outArry = []
 outVid = cv2.VideoWriter(args.save_stem + '.avi', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (OUTW, OUTH))
 
@@ -141,4 +153,7 @@ bar.finish()
 print(f'Ending at time = {time}')
 
 # Save out array
-np.savez_compressed(args.save_stem + '.npz', np.array(outArry))
+#np.savez_compressed(args.save_stem + '.npz', np.array(outArry))
+with h5py.File(args.save_stem + '.hdf5', 'w') as f:
+    f.create_dataset('data', data=np.array(outArry))
+    f.attrs['metadata'] = json.dumps(metadata)