From 4ec37b4164477e7868ee251d8a38ff3c39cae98e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Agust=C3=ADn=20Castro?= <acastro@192.168.1.110>
Date: Thu, 8 Feb 2024 11:43:08 -0300
Subject: [PATCH] Fix output resolution error and slight disalignment of frames

---
 demos/multi_camera/demo.py | 59 ++++++++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 15 deletions(-)

diff --git a/demos/multi_camera/demo.py b/demos/multi_camera/demo.py
index e257b88b..a68a1d69 100644
--- a/demos/multi_camera/demo.py
+++ b/demos/multi_camera/demo.py
@@ -1,7 +1,6 @@
 import argparse
 import os
 import pickle
-import sys
 
 import cv2
 import numpy as np
@@ -353,8 +352,22 @@ def mask_generator(frame):
         if args.use_motion_estimator_footage:
             motion_estimators[path].transformation = initial_transformations[path]
 
-    # now initialize the videos and their trackers
+    # initialize the reference if it exists
+    reference = {"video": None, "image": None, "motion_estimator": None}
     image_reference = None
+    if not first_video_is_reference:
+        # if failing to read it as an image, try to read it as a video
+        image_reference = cv2.imread(args.reference)
+        reference["image"] = image_reference
+        if image_reference is None:
+            video = Video(input_path=path)
+            image_reference = next(video.__iter__())
+            reference["video"] = video
+            reference["motion_estimator"] = motion_estimator_reference
+
+    # now initialize the videos and their trackers
+    fps = None
+    total_frames = None
     for path in args.files:
         extension = os.path.splitext(path)[1]
         if args.output_name is None:
@@ -363,8 +376,26 @@ def mask_generator(frame):
             output_path = args.output_name
 
         video = Video(input_path=path, output_path=output_path)
+        if fps is None:
+            fps = video.output_fps
+            total_frames = int(video.video_capture.get(cv2.CAP_PROP_FRAME_COUNT))
+        else:
+            current_fps = video.output_fps
+            current_total_frames = int(
+                video.video_capture.get(cv2.CAP_PROP_FRAME_COUNT)
+            )
+            if current_fps != fps:
+                raise ValueError(
+                    f"{args.files[0]} is at {fps} FPS, but {path} is at {current_fps} FPS."
+                )
+            if total_frames != current_total_frames:
+                raise ValueError(
+                    f"{args.files[0]} has {total_frames} frames, but {path} has {current_total_frames} frames."
+                )
         if image_reference is None:
             image_reference = next(video.__iter__())
+        else:
+            next(video.__iter__())
         videos[path] = video
         trackers[path] = Tracker(
             distance_function="iou",
@@ -376,22 +407,20 @@ def mask_generator(frame):
         )
         tracked_objects[path] = []
 
-    reference = {"video": None, "image": None, "motion_estimator": None}
-    if not first_video_is_reference:
-        # if failing to read it as an image, try to read it as a video
-        image_reference = cv2.imread(args.reference)
-        reference["image"] = image_reference
-        if image_reference is None:
-            video = Video(input_path=path)
-            image_reference = next(video.__iter__())
-            reference["video"] = video
-            reference["motion_estimator"] = motion_estimator_reference
-
-    height_reference = image_reference.shape[0]
     big_black_frame = np.zeros(
-        tuple([args.resolution[1] * 2, args.resolution[0] * 2, 3]), dtype=np.uint8
+        tuple(
+            [
+                args.resolution[1]
+                * ((len(args.files) + (not first_video_is_reference) + 1) // 2),
+                args.resolution[0] * 2,
+                3,
+            ]
+        ),
+        dtype=np.uint8,
     )
 
+    height_reference = image_reference.shape[0]
+
     def normalized_foot_distance(tracker1, tracker2):
         return (
             np.linalg.norm(get_absolute_feet(tracker1) - get_absolute_feet(tracker2))