with a coffee

dodgeballtracker/__main__.py

@@ -6,8 +6,7 @@ from dodgeballtracker.metadata import VideoMetadata
 
 
 from dodgeballtracker.helper import time_tracker
-from dodgeballtracker.logic import (DodgeballTracker,
-                                    get_court_perspective_transform)
+from dodgeballtracker.logic import DodgeballTracker
 from dodgeballtracker.data_dumper import SQLiteDataDumper, DataDumper
 import numpy as np
 import logging
@@ -103,7 +102,7 @@ def train():
 @click.option('--input', '-i', 'input_',
               type=click.Path(readable=True, path_type=Path))
 @click.option('--output-dir', '-o', type=click.Path(path_type=Path))
-@click.option('--metadata-file', '-f', type=click.Path(path_type=Path, file_okay=True, exists=True))
+@click.option('--metadata-file', '-m', type=click.Path(path_type=Path, file_okay=True, exists=True))
 @click.option('--display-video', '-d', type=bool,  is_flag=True, show_default=True, default=False)
 def process(input_: Path,
             output_dir: Path | None,

dodgeballtracker/data_dumper.py

@@ -13,7 +13,9 @@ import pandera as pa
 class DataDumper(abc.ABC):
 
     @abc.abstractmethod
-    def write_boxes(self, obj: list[TrackedDetectedObject], ts: VideoTimestamp,
+    def write_boxes(self,
+                    tracked_obj: pd.DataFrame,
+                    ts: VideoTimestamp,
                     frame_index: int) -> None:
         pass
 
@@ -78,7 +80,7 @@ class SQLiteDataDumper(DataDumper):
     def start(self,
               video_path: Path,
               hexa_hash: str,
-              metadata: VideoMetadata|None=None) -> None:
+              metadata: VideoMetadata | None = None) -> None:
         self._create_tables()
 
         cur = self._conn.cursor()
@@ -110,7 +112,8 @@ class SQLiteDataDumper(DataDumper):
         RETURNING id
         ''', (ts, frame_index, self._video_id))
         frame_id = res.fetchone()[0]
-        cols = ['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category', 'object_id']
+        cols = ['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category',
+                'object_id']
         cur.executemany('''
             INSERT INTO tracked_detected_object(xmin, 
                                                 ymin,

dodgeballtracker/logic.py

@@ -28,31 +28,21 @@ def get_best_available_torch_device() -> torch.device:
     return torch.device('cpu')
 
 
-def get_court_perspective_transform():
-    """
-    """
-    empirical_top_left = np.float32([416, 175])
-    empirical_top_right = np.float32([905, 173])
-    empirical_bottom_left = np.float32([51, 702])
-    empirical_bottom_right = np.float32([1188, 695])
-
-    court_top_left = np.float32([0, 0])
-    court_top_right = np.float32([9, 0])
-    court_bottom_left = np.float32([0, 18])
-    court_bottom_right = np.float32([9, 18])
-
-    empirical_source = np.float32([empirical_top_left,
-                                   empirical_top_right,
-                                   empirical_bottom_left,
-                                   empirical_bottom_right])
-    pixel_per_meter = 50
-    empirical_target = np.float32([court_top_left,
-                                   court_top_right,
-                                   court_bottom_left,
-                                   court_bottom_right]) * pixel_per_meter
-    return cv2.getPerspectiveTransform(empirical_source, empirical_target)
 
+class BirdViewCreator:
+    def __init__(self, view_perspective_matrix: np.matrix):
+        self._view_perspective_matrix = view_perspective_matrix
 
+    @pa.check_io(df=TrackedDetectedObject, out=TrackedDetectedObject)
+    def append_bird_view(self, df: pd.DataFrame) -> pd.DataFrame:
+        v = np.matrix([
+            pd.Series(dtype='float32', data=np.ones(df.shape[0])),
+            df['xmin'],
+            df['xmax']
+        ])
+        ret = self._view_perspective_matrix @ v
+        ret_df = pd.DataFrame(ret).rename(columns={0: 'physical_x', 1: 'physical_y'})
+        return df.join(ret_df)
 
 
 
@@ -63,12 +53,12 @@ class DodgeballTracker:
                  video_writer: VideoWriter,
                  data_dumper: DataDumper,
                  detected_box_objects: list[DetectedObjectBox],
-                 transformation_matrix):
+                 bird_view_creator: BirdViewCreator):
         self.input_path = input_file
         self.video_writer = video_writer
         self._data_dumper = data_dumper
         self.detected_box_objects = sorted(detected_box_objects, reverse=True)
-        self.transformation_matrix = transformation_matrix
+        self._bird_view_creator = bird_view_creator
 
 
     @pa.check_output(DetectedObject)
@@ -103,9 +93,6 @@ class DodgeballTracker:
             return pd.DataFrame({col: pd.Series(dtype='float64') for col in columns})
         return pd.DataFrame(data, columns=columns)
 
-    def train(self):
-        model = self._get_detection_model()
-
     def main(self):
         # define some constants
         video_cap = cv2.VideoCapture(str(self.input_path))
@@ -116,7 +103,10 @@ class DodgeballTracker:
             # run the YOLO model on the frame
             detected_objs = self.run_model_on_frame(model, frame)
             tracked_obj = self.update_tracker(tracker, detected_objs, frame)
-            self._data_dumper.write_boxes(tracked_obj, ts, i)
+            self._data_dumper.write_boxes(
+                self._bird_view_creator.append_bird_view(tracked_obj),
+                ts,
+                i)
             self.video_writer.plot_boxes_to_frame(tracked_obj, frame)
             self.video_writer.write_frame(frame)
 
@@ -124,3 +114,5 @@ class DodgeballTracker:
                 break
 
         video_cap.release()
+
+

dodgeballtracker/metadata.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from pydantic import BaseModel
 import numpy as np
 import cv2
@@ -31,8 +29,8 @@ class KnownPoint(BaseModel):
 class VideoMetadata(BaseModel):
     knownPoints: list[KnownPoint]
 
-    def get_view_perspective_matrix(self) -> Optional[np.array]:
-        if len(self.knownPoints) < 4:
+    def get_view_perspective_matrix(self) -> np.ndarray | None:
+        if len(self.knownPoints) != 4:
             return None
 
         empirical_source = np.float32(

dodgeballtracker/typing.py

@@ -16,4 +16,5 @@ TrackedDetectedObject = tensor_of(Float64,
                                   list(DetectedObject.columns) + ["object_id", "index"])
 
 
+
 Color = tuple[int, int, int]