with a glass of water

.gitignore

@@ -164,3 +164,4 @@ result
 dodgeballtracker/runs
 data
 outputs
+runs

dodgeballtracker/__main__.py

@@ -82,15 +82,13 @@ def process_logic(input_: Path,
         np.arange(0.5, 1, 0.05)
     ]
 
-    transformation_matrix = get_court_perspective_transform()
-
     with get_video_writer(display_video, input_, output_dir / 'video.mov') as video_writer:
         with create_data_dumper(input_, output_dir / 'data.sqlite', metadata) as data_dumper:
             tracker = DodgeballTracker(input_,
                                        video_writer,
                                        data_dumper,
                                        detected_box_objects,
-                                       transformation_matrix)
+                                       metadata.get_view_perspective_matrix())
 
             with time_tracker():
                 tracker.main()

dodgeballtracker/data_dumper.py

@@ -61,12 +61,14 @@ class SQLiteDataDumper(DataDumper):
                                     ymin DOUBLE PRECISION NOT NULL,
                                     xmax DOUBLE PRECISION NOT NULL,
                                     ymax DOUBLE PRECISION NOT NULL,
+                                    
+                                    bird_view_x DOUBLE PRECISION,
+                                    bird_view_y DOUBLE PRECISION,
+                                    
                                     score DOUBLE PRECISION NOT NULL,
                                     category INTEGER NOT NULL,
                                     object_id INTEGER NOT NULL,
                                     video_frame_id INTEGER NOT NULL,
-                                    
-                                    
                                     FOREIGN KEY(video_frame_id) REFERENCES video_frame(id)
                                 );
                                 ''')

dodgeballtracker/logic.py

@@ -1,4 +1,3 @@
-import json
 from pathlib import Path
 import logging
 import torch
@@ -70,7 +69,6 @@ class DodgeballTracker:
         self._data_dumper = data_dumper
         self.detected_box_objects = sorted(detected_box_objects, reverse=True)
         self.transformation_matrix = transformation_matrix
-        self.points = []
 
 
     @pa.check_output(DetectedObject)
@@ -94,7 +92,7 @@ class DodgeballTracker:
         load the pre-trained YOLOv8n model
         """
         my_model = Path(
-            '/home/didier/all-dressed-programming/dodgeball-tracker/dodgeballtracker/runs/detect/train31/weights/best.pt')
+            '/home/didier/all-dressed-programming/dodgeball-tracker/runs/detect/train31/weights/best.pt')
         return YOLO(my_model, verbose=False)
 
     @pa.check_io(detected_objs=DetectedObject, out=TrackedDetectedObject)

dodgeballtracker/metadata.py

@@ -1,14 +1,27 @@
+from typing import Optional
+
 from pydantic import BaseModel
+import numpy as np
+import cv2
+
 
 class KnownImagePoint(BaseModel):
-    height: float
     width: float
+    height: float
+
+    @classmethod
+    def from_vec(cls, vec: tuple[float, float]) -> 'KnownImagePoint':
+        return cls(width=vec[0], height=vec[1])
 
 
 class BirdViewPoint(BaseModel):
     x: float
     y: float
 
+    @classmethod
+    def from_vec(cls, vec: tuple[float, float]) -> 'BirdViewPoint':
+        return cls(x=vec[0], y=vec[1])
+
 
 class KnownPoint(BaseModel):
     image: KnownImagePoint
@@ -18,3 +31,14 @@ class KnownPoint(BaseModel):
 class VideoMetadata(BaseModel):
     knownPoints: list[KnownPoint]
 
+    def get_view_perspective_matrix(self) -> Optional[np.array]:
+        if len(self.knownPoints) < 4:
+            return None
+
+        empirical_source = np.float32(
+            [[k.image.width, k.image.height] for k in self.knownPoints])
+        empirical_target = np.float32(
+            [[k.bird_view.x, k.bird_view.y] for k in self.knownPoints])
+        pixel_per_meter = 100
+        return cv2.getPerspectiveTransform(empirical_source,
+                                           empirical_target * pixel_per_meter)

tests/test_metadata.py

@@ -0,0 +1,83 @@
+import pytest
+import numpy as np
+from dodgeballtracker.metadata import (VideoMetadata, KnownImagePoint,
+                                       BirdViewPoint, KnownPoint)
+
+
+@pytest.fixture
+def simple_transformation_matrix() -> np.matrix:
+    m = np.matrix([[2, 1], [1, 1]])
+    assert np.linalg.det(m) > 0.1
+    return m
+
+
+def create_know_point(v, m) -> KnownPoint:
+    u = (m @ v).A[0]
+    return KnownPoint(image=KnownImagePoint.from_vec(v),
+                      bird_view=BirdViewPoint.from_vec(u))
+
+
+@pytest.fixture
+def four_known_points(simple_transformation_matrix: np.matrix) -> list[
+    KnownPoint]:
+    return [
+        create_know_point([x, y], simple_transformation_matrix)
+        for x in [-100, 100]
+        for y in [-100, 100]
+    ]
+
+
+@pytest.fixture
+def eight_known_points(simple_transformation_matrix: np.matrix,
+                       four_known_points: list[KnownPoint]) -> list[KnownPoint]:
+    return [
+        create_know_point([x, y], simple_transformation_matrix)
+        for x in [-50, 50]
+        for y in [-50, 50]
+    ] + four_known_points
+
+
+
+@pytest.fixture
+def simple_video_metadata(four_known_points) -> VideoMetadata:
+    return VideoMetadata(knownPoints=four_known_points)
+
+@pytest.fixture
+def mega_video_metadata(eight_known_points) -> VideoMetadata:
+    return VideoMetadata(knownPoints=eight_known_points)
+
+@pytest.fixture
+def broken_video_metadata(four_known_points) -> VideoMetadata:
+    return VideoMetadata(knownPoints=four_known_points[:2])
+
+
+def test_get_view_perspective_matrix(simple_video_metadata: VideoMetadata) -> None:
+    m = simple_video_metadata.get_view_perspective_matrix()
+    expt = np.array([
+        [200, 100, 0],
+        [100, 100, 0],
+        [0, 0, 1]
+    ])
+    np.testing.assert_array_almost_equal(m, expt)
+
+
+@pytest.mark.xfail
+def test_get_view_perspective_matrix_mega(mega_video_metadata: VideoMetadata) -> None:
+    """
+    The issue with this test is that the transformation
+    matrix can only be computed is there's exactly 4 known
+    points.  Obviously, in the real world we want more known
+    points and find a way to get a better estimate.
+    """
+    m = mega_video_metadata.get_view_perspective_matrix()
+    expt = np.array([
+        [200, 100, 0],
+        [100, 100, 0],
+        [0, 0, 1]
+    ])
+    np.testing.assert_array_almost_equal(m, expt)
+
+
+def test_get_view_perspective_matrix_without_enough_data(broken_video_metadata: VideoMetadata) -> None:
+    m = broken_video_metadata.get_view_perspective_matrix()
+    assert m is None