with a coffee

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()
+
+