From 31b9e85912a204d1130453b19b099381cce2ba3e Mon Sep 17 00:00:00 2001 From: didier amyot Date: Sat, 16 Nov 2024 20:35:27 -0500 Subject: [PATCH] with a glass of water --- .gitignore | 1 + dodgeballtracker/__main__.py | 4 +- dodgeballtracker/data_dumper.py | 6 ++- dodgeballtracker/logic.py | 4 +- dodgeballtracker/metadata.py | 26 ++++++++++- tests/test_metadata.py | 83 +++++++++++++++++++++++++++++++++ 6 files changed, 115 insertions(+), 9 deletions(-) create mode 100644 tests/test_metadata.py diff --git a/.gitignore b/.gitignore index de69d12..56b0e08 100644 --- a/.gitignore +++ b/.gitignore @@ -164,3 +164,4 @@ result dodgeballtracker/runs data outputs +runs diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index 42bae72..5c75a5e 100644 --- a/dodgeballtracker/__main__.py +++ b/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() diff --git a/dodgeballtracker/data_dumper.py b/dodgeballtracker/data_dumper.py index 090be17..37e83cb 100644 --- a/dodgeballtracker/data_dumper.py +++ b/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) ); ''') diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 7a9d5e8..5073d43 100644 --- a/dodgeballtracker/logic.py +++ b/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) diff --git a/dodgeballtracker/metadata.py b/dodgeballtracker/metadata.py index 7638805..f233ff2 100644 --- a/dodgeballtracker/metadata.py +++ b/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) diff --git a/tests/test_metadata.py b/tests/test_metadata.py new file mode 100644 index 0000000..5727766 --- /dev/null +++ b/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