.envrc

@@ -0,0 +1 @@
+use flake

.gitignore

@@ -159,3 +159,4 @@ env
 *.pt
 *.png
 *.ipynb
+.direnv

dodgeballtracker/__main__.py

@@ -1,28 +1,69 @@
+import time
+from contextlib import closing
 from pathlib import Path
 
+from dodgeballtracker.helper import time_tracker
 from dodgeballtracker.logic import (DetectedObjectBox,
                                     DodgeballTracker,
                                     get_court_perspective_transform)
 import numpy as np
 import logging
+import click
+from typing import ContextManager
 
+from dodgeballtracker.video_writer import (MP4VVideoWriter, VideoWriter,
+                                           MP4VVideoWriterWithFrame)
 
 here = Path(__file__).parent
 input_dir = here / '..' / 'data'
 output_dir = here / '..' / 'outputs'
 
-if __name__ == "__main__":
+
+@click.group()
+@click.option('--debug/--no-debug', default=False)
+def cli(debug):
     logging.basicConfig(filename='tracker.log', level=logging.DEBUG)
 
-    detected_box_objects = [
-        DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05)
-    ]
 
-    input_video = input_dir / "small.mp4"
+def get_output(input_: Path, proposed_output: Path | None) -> Path:
-    output_video = output_dir / "foo.mp4"
+    if proposed_output is not None:
+        return proposed_output
+    return output_dir / input_.name
+
+
+def get_video_writer(display_frame: bool,
+                     video_input: Path,
+                     video_output: Path) -> ContextManager[VideoWriter]:
+    if display_frame:
+        return closing(
+            MP4VVideoWriterWithFrame.from_frame(video_input, video_output))
+    return closing(MP4VVideoWriter.from_frame(video_input, video_output))
+
+
+@cli.command()
+@click.option('--input', '-i', 'input_',
+              type=click.Path(readable=True, path_type=Path))
+@click.option('--output', '-o', type=click.Path(path_type=Path))
+@click.option('--display-video', '-d', type=bool,  is_flag=True, show_default=True, default=False)
+def process(input_: Path, output: Path | None, display_video: bool):
+    output = get_output(input_, output)
+    detected_box_objects = [
+        DetectedObjectBox.from_viridis(threshold) for threshold in
+        np.arange(0.5, 1, 0.05)
+    ]
 
     transformation_matrix = get_court_perspective_transform()
 
-    dodgeball_tracker = DodgeballTracker(input_video, output_video, detected_box_objects, transformation_matrix)
+    with get_video_writer(display_video, input_, output) as video_writer:
+        tracker = DodgeballTracker(input_,
+                                   output,
+                                   video_writer,
+                                   detected_box_objects,
+                                   transformation_matrix)
 
-    dodgeball_tracker.main()
+        with time_tracker():
+            tracker.main()
+
+
+if __name__ == "__main__":
+    cli()

dodgeballtracker/helper.py

@@ -1,21 +1,10 @@
 import contextlib
 import time
 import cv2
+import logging
 
+logger = logging.getLogger(__name__)
 
-def create_video_writer(video_cap, output_filename):
-
-    # grab the width, height, and fps of the frames in the video stream.
-    frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    fps = int(video_cap.get(cv2.CAP_PROP_FPS))
-
-    # initialize the FourCC and a video writer object
-    fourcc = cv2.VideoWriter_fourcc(*'MP4V')
-    writer = cv2.VideoWriter(output_filename, fourcc, fps,
-                             (frame_width, frame_height))
-
-    return writer
 
 
 def iter_frames(video_cap):
@@ -31,4 +20,4 @@ def time_tracker():
     start = time.monotonic()
     yield
     end = time.monotonic()
-    print(f"Time to process 1 frame: {(end-start) * 1000:.0f} milliseconds")
+    logger.debug(f"Time to process 1 frame: {(end-start) * 1000:.0f} milliseconds")

dodgeballtracker/logic.py

@@ -4,7 +4,8 @@ import logging
 import torch
 from ultralytics import YOLO
 import cv2
-from dodgeballtracker.helper import create_video_writer, iter_frames
+from dodgeballtracker.helper import iter_frames
+from dodgeballtracker.video_writer import create_video_writer, VideoWriter
 import matplotlib
 import numpy as np
 from boxmot import DeepOcSort as DeepOCSORT
@@ -15,15 +16,15 @@ logger = logging.getLogger(__name__)
 Color = tuple[int, int, int]
 
 
-def get_best_available_torch_device():
+def get_best_available_torch_device() -> torch.device:
     if torch.cuda.is_available():
         logger.debug("Using the cuda (gpu) torch device.")
-        return 'cuda'
+        return torch.device('cuda:0')
     if torch.backends.mps.is_available():
         logger.debug("Using the mps (apple silicon) torch device.")
-        return 'mps'
+        return torch.device('mps')
     logger.debug("Default to cpu (no apple silicon or cuda) torch device.")
-    return 'cpu'
+    return torch.device('cpu')
 
 
 class DetectedObjectBox:
@@ -66,16 +67,21 @@ def get_court_perspective_transform():
     return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
 
+
+
+
 class DodgeballTracker:
 
     def __init__(self,
                  input_file: Path,
                  output_file: Path,
+                 video_writer: VideoWriter,
                  detected_box_objects: list[DetectedObjectBox],
                  transformation_matrix):
         self.input_path = input_file
         self.output_path = output_file
         self.output_json_path = output_file.with_suffix(".json")
+        self.video_writer = video_writer
         self.detected_box_objects = sorted(detected_box_objects, reverse=True)
         self.transformation_matrix = transformation_matrix
         self.points = []
@@ -138,7 +144,7 @@ class DodgeballTracker:
     def _get_tracker(self):
         return DeepOCSORT(
             Path('lmbn_n_cuhk03_d.pt'),  # which ReID model to use
-            torch.device(get_best_available_torch_device()),
+            get_best_available_torch_device(), # torch.cuda.get_device_name(0), #torch.device(get_best_available_torch_device()),
             True,
         )
 
@@ -150,10 +156,8 @@ class DodgeballTracker:
 
     def main(self):
         # define some constants
-        print("started")
-
         video_cap = cv2.VideoCapture(str(self.input_path))
-        writer = create_video_writer(video_cap, str(self.output_path))
+        # writer = create_video_writer(video_cap, self.output_path)
         tracker = self._get_tracker()
 
         # load the pre-trained YOLOv8n model
@@ -168,19 +172,13 @@ class DodgeballTracker:
                 if is_good_box:
                     res.append(data)
 
-            # show the frame to our screen
-            print("tracker?")
             data_ = tracker.update(np.array(res), frame)
             self.print_tracker_data(data_, i, ts)
 
-            self.write_frame(writer, frame)
+            self.video_writer.write_frame(frame)
 
             if cv2.waitKey(1) == ord("q"):
                 break
-        print("fini")
 
         video_cap.release()
-        writer.release()
-        cv2.destroyAllWindows()
         self.persist_points()
-        print("finished")

dodgeballtracker/video_writer.py

@@ -0,0 +1,67 @@
+import abc
+from pathlib import Path
+
+import cv2
+
+def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter:
+
+    # grab the width, height, and fps of the frames in the video stream.
+    frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    fps = int(video_cap.get(cv2.CAP_PROP_FPS))
+
+    # initialize the FourCC and a video writer object
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    writer = cv2.VideoWriter(str(output_filename.absolute()),
+                             fourcc,
+                             fps,
+                             (frame_width, frame_height))
+
+    return writer
+
+
+class VideoWriter(abc.ABC):
+
+    @abc.abstractmethod
+    def write_frame(self, frame) -> None:
+        pass
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        pass
+
+
+class NoopVideoWriter(VideoWriter):
+
+    def write_frame(self, frame) -> None:
+        pass
+
+    def close(self) -> None:
+        pass
+
+
+class MP4VVideoWriter(VideoWriter):
+
+    def __init__(self, video_writer: cv2.VideoWriter) -> None:
+        self._video_writer = video_writer
+
+    def write_frame(self, frame) -> None:
+        self._video_writer.write(frame)
+
+    def close(self) -> None:
+        self._video_writer.release()
+        cv2.destroyAllWindows()
+
+    @classmethod
+    def from_frame(cls, input_: Path, output: Path) -> VideoWriter:
+        video_cap = cv2.VideoCapture(str(input_))
+        try:
+            return cls(create_video_writer(video_cap, output))
+        finally:
+            video_cap.release()
+
+
+class MP4VVideoWriterWithFrame(MP4VVideoWriter):
+    def write_frame(self, frame) -> None:
+        super().write_frame(frame)
+        cv2.imshow("Frame", frame)

flake.nix

@@ -27,9 +27,10 @@
         virtual-env = pkgs.python312.withPackages (ps: [
           ps.matplotlib
           ps.torchWithCuda
-          # pythonPackages.opencv4
+          pythonPackages.opencv4
-          # pythonPackages.boxmot
+          pythonPackages.boxmot
-          # pythonPackages.ultralytics
+          pythonPackages.ultralytics
+          pythonPackages.click
         ]);
 
       in {