Fix video 'pred_only' issue

2024-06-20 05:23:51 -07:00
parent 3f5fe63b59
commit 424bc41a26
2 changed files with 6 additions and 2 deletions
--- a/run_video.py
+++ b/run_video.py
@@ -56,7 +56,11 @@ if __name__ == '__main__':
        raw_video = cv2.VideoCapture(filename)
        frame_width, frame_height = int(raw_video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(raw_video.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frame_rate = int(raw_video.get(cv2.CAP_PROP_FPS))
-        output_width = frame_width * 2 + margin_width
+        
+        if args.pred_only: 
+            output_width = frame_width
+        else: 
+            output_width = frame_width * 2 + margin_width
        
        output_path = os.path.join(args.outdir, os.path.splitext(os.path.basename(filename))[0] + '.mp4')
        out = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*"mp4v"), frame_rate, (output_width, frame_height))