From 0f0fe0f518980088e380e4dd4cabb34d263c1175 Mon Sep 17 00:00:00 2001 From: didier amyot Date: Sat, 23 Nov 2024 09:16:24 -0500 Subject: [PATCH] with a coffee --- dodgeballtracker/__main__.py | 5 ++-- dodgeballtracker/data_dumper.py | 9 ++++-- dodgeballtracker/logic.py | 50 ++++++++++++++------------------- dodgeballtracker/metadata.py | 6 ++-- dodgeballtracker/typing.py | 1 + tests/test_data_dumper.py | 5 ++-- tests/test_logic.py | 39 +++++++++++++++++++++++++ 7 files changed, 74 insertions(+), 41 deletions(-) create mode 100644 tests/test_logic.py diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index 5c75a5e..9a1e248 100644 --- a/dodgeballtracker/__main__.py +++ b/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, diff --git a/dodgeballtracker/data_dumper.py b/dodgeballtracker/data_dumper.py index 37e83cb..1fe1de3 100644 --- a/dodgeballtracker/data_dumper.py +++ b/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, diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 5073d43..d7f491c 100644 --- a/dodgeballtracker/logic.py +++ b/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() + + diff --git a/dodgeballtracker/metadata.py b/dodgeballtracker/metadata.py index f233ff2..3dfb614 100644 --- a/dodgeballtracker/metadata.py +++ b/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( diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index 78a8e3c..170dc0f 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -16,4 +16,5 @@ TrackedDetectedObject = tensor_of(Float64, list(DetectedObject.columns) + ["object_id", "index"]) + Color = tuple[int, int, int] diff --git a/tests/test_data_dumper.py b/tests/test_data_dumper.py index a7eb457..2ef1b53 100644 --- a/tests/test_data_dumper.py +++ b/tests/test_data_dumper.py @@ -98,7 +98,8 @@ def test_restart(tmp_dir: Path): assert expt == list_tables(conn) -def test_write_boxes_simple(tmp_dir: Path, mock_tracked_detected_object: pd.DataFrame): +def test_write_boxes_simple(tmp_dir: Path, + mock_tracked_detected_object: pd.DataFrame): with make_conn(tmp_dir) as conn: dumper = SQLiteDataDumper(conn) dumper.start(Path('/foo/toto.MOV'), 'aabbcc') @@ -130,7 +131,7 @@ def test_write_boxes_with_empty_frame(tmp_dir: Path, with make_conn(tmp_dir) as conn: dumper = SQLiteDataDumper(conn) dumper.start(Path('/foo/toto.MOV'), 'aabbcc') - dumper.write_boxes(empty_tracked_detected_object, 40018, 34) + dumper.write_boxes(empty_tracked_detected_object, None,40018, 34) with make_conn(tmp_dir) as conn: assert 0 == number_of_items_in_table('tracked_detected_object', conn) diff --git a/tests/test_logic.py b/tests/test_logic.py new file mode 100644 index 0000000..151e246 --- /dev/null +++ b/tests/test_logic.py @@ -0,0 +1,39 @@ +import numpy as np +import pandas as pd +import pandera as pa +import pytest + +from dodgeballtracker.logic import BirdViewCreator +from dodgeballtracker.typing import TrackedDetectedObject + + +@pytest.fixture +def view_perspective() -> np.matrix: + return np.matrix([ + [200, 100, 0], + [100, 100, 0], + [0, 0, 1] + ]) + + +@pytest.fixture +@pa.check_io(out=TrackedDetectedObject) +def mock_data() -> pd.DataFrame: + return pd.DataFrame( + { + "xmin": pd.Series(dtype='float64', data=[10.0, 100.0]), + "ymin": pd.Series(dtype='float64', data=[30.0, 130.0]), + "xmax": pd.Series(dtype='float64', data=[55.0, 150.0]), + "ymax": pd.Series(dtype='float64', data=[65.0, 195.0]), + "score": pd.Series(dtype='float64', data=[0.66, 0.56]), + "category": pd.Series(dtype='float64', data=[3, 3]), + "object_id": pd.Series(dtype='float64', data=[73232, 837276]), + "index": pd.Series(dtype='float64', data=[0, 1]), + } + ) + + +def test_get_bird_view(view_perspective: np.matrix, + mock_data: pd.DataFrame) -> None: + bvd = BirdViewCreator(view_perspective) + ret = bvd.append_bird_view(mock_data)