Fix demo CLI

Author: cjaverliat

kineo/demo/offline/demo.py

@@ -71,8 +71,7 @@ def main(config_file: str, sequence_name: str, video_paths: list[str], batch_siz
         gt_annotations={}
     )
 
-
-if __name__ == "__main__":
+def cli():
     parser = argparse.ArgumentParser()
     parser.add_argument("--config-file", type=str, default="configs/demo/offline/nlf_single_person_sam2.yaml")
     parser.add_argument("--sequence-name", type=str, default="offline_demo")
@@ -82,11 +81,16 @@ def main(config_file: str, sequence_name: str, video_paths: list[str], batch_siz
     parser.add_argument("--use-cache", action="store_true", default=False)
     parser.add_argument("video_paths", type=str, nargs="+")
     args = parser.parse_args()
-    config_file = args.config_file
-    sequence_name = args.sequence_name
-    video_paths = args.video_paths
-    batch_size = args.batch_size
-    target_fps = args.target_fps
-    shared_intrinsics = args.shared_intrinsics
-    use_cache = args.use_cache
-    main(config_file, sequence_name, video_paths, batch_size, target_fps, shared_intrinsics, use_cache)
+
+    main(
+        args.config_file,
+        args.sequence_name,
+        args.video_paths,
+        args.batch_size,
+        args.target_fps,
+        args.shared_intrinsics,
+        args.use_cache,
+    )
+
+if __name__ == "__main__":
+    cli()

kineo/demo/online/demo.py

@@ -21,7 +21,7 @@
 torch.backends.cudnn.benchmark = True
 
 
-def create_views_from_temp_videos(temp_videos, cam_names):
+def create_views_from_temp_videos(temp_videos, cam_names, device: torch.device):
     views = []
     for i, temp_video in enumerate(temp_videos):
         video_path = temp_video.name
@@ -38,7 +38,7 @@ def create_views_from_temp_videos(temp_videos, cam_names):
         )
     return views
 
-def create_live_views(cam_indices: list[int], cam_names: list[str]):
+def create_live_views(cam_indices: list[int], cam_names: list[str], device: torch.device):
     views = []
     for cam_idx, cam_name in zip(cam_indices, cam_names):
         views.append(
@@ -82,43 +82,39 @@ def load_camera_calibrations(cam_ids: list[str], calibration_output_root_dir: st
 
     return cam_intrinsics, cam_extrinsics, world_reconstructed_scene
 
-if __name__ == "__main__":
 
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--target-fps", type=int, default=20)
-    parser.add_argument("--target-res", type=str, default="640x480")
-    parser.add_argument("--live-viz-config", type=str, default="configs/demo/realtime/realtime_viz.yaml")
-    parser.add_argument("--skip-calibration", action="store_true")
-    args = parser.parse_args()
-    target_fps = int(args.target_fps)
-    target_res = tuple(int(x) for x in args.target_res.split("x"))
+def main(target_fps, target_res, live_viz_config, skip_calibration):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
     calibration_config_file = "configs/demo/realtime/calibration.yaml"
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    if not args.skip_calibration:
+    if not skip_calibration:
         print("Loading calibration pipeline")
         calibration_pipeline = Pipeline.build_pipeline_from_config(calibration_config_file, device)
         print("Pipelines loaded")
 
         app = QtWidgets.QApplication(sys.argv)
-        recorder = MultiCamRecorder(target_fps=target_fps, target_res=target_res, api_preference=cv2.CAP_ANY)
+        recorder = MultiCamRecorder(
+            target_fps=target_fps,
+            target_res=target_res,
+            api_preference=cv2.CAP_ANY
+        )
         recorder.show()
         app.exec_()
 
         cam_indices = recorder.camera_indices
         cam_names = recorder.camera_ids
 
-        views = create_views_from_temp_videos(recorder.temp_videos, cam_names)
-        annotations = calibration_pipeline.run(
+        views = create_views_from_temp_videos(recorder.temp_videos, cam_names, device=device)
+        calibration_pipeline.run(
             sequence_name="calibration",
             views=views,
             annotations={},
             gt_annotations={},
         )
 
     print("Loading live viz pipeline")
-    live_viz_pipeline = Pipeline.build_pipeline_from_config(args.live_viz_config, device)
+    live_viz_pipeline = Pipeline.build_pipeline_from_config(live_viz_config, device)
     print("Live viz pipeline loaded")
 
     camera_indices = []
@@ -132,13 +128,14 @@ def load_camera_calibrations(cam_ids: list[str], calibration_output_root_dir: st
     if len(camera_indices) <= 1:
         raise Exception(f"Expected at least 2 cameras, got {len(camera_indices)}")
 
-    views = create_live_views(camera_indices, camera_ids)
+    views = create_live_views(camera_indices, camera_ids, device=device)
 
-    # Load calibration data
     calibration_output_root_dir = "./outputs/realtime_demo_calibration"
-    cam_intrinsics, cam_extrinsics, world_reconstructed_scene = load_camera_calibrations(camera_ids, calibration_output_root_dir)
+    cam_intrinsics, cam_extrinsics, world_reconstructed_scene = load_camera_calibrations(
+        camera_ids, calibration_output_root_dir
+    )
 
-    _ = live_viz_pipeline.run(
+    live_viz_pipeline.run(
         sequence_name="realtime_viz",
         views=views,
         annotations={
@@ -147,4 +144,25 @@ def load_camera_calibrations(cam_ids: list[str], calibration_output_root_dir: st
             "world_reconstructed_scene": world_reconstructed_scene,
         },
         gt_annotations={},
-    )
+    )
+
+
+def cli():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--target-fps", type=int, default=20)
+    parser.add_argument("--target-res", type=str, default="640x480")
+    parser.add_argument("--live-viz-config", type=str, default="configs/demo/realtime/realtime_viz.yaml")
+    parser.add_argument("--skip-calibration", action="store_true")
+    args = parser.parse_args()
+
+    target_res = tuple(int(x) for x in args.target_res.split("x"))
+
+    main(
+        target_fps=args.target_fps,
+        target_res=target_res,
+        live_viz_config=args.live_viz_config,
+        skip_calibration=args.skip_calibration,
+    )
+
+if __name__ == "__main__":
+    cli()

pyproject.toml

@@ -53,8 +53,8 @@ dev = [
 ]
 
 [project.scripts]
-kineo-offline = "kineo.demo.offline.demo:main"
-kineo-online = "kineo.demo.online.demo:main"
+kineo-offline = "kineo.demo.offline.demo:cli"
+kineo-online = "kineo.demo.online.demo:cli"
 
 [tool.setuptools.packages.find]
 where = ["."]