diff --git a/.gitignore b/.gitignore index a56370a..de69d12 100644 --- a/.gitignore +++ b/.gitignore @@ -161,3 +161,6 @@ env *.ipynb .direnv result +dodgeballtracker/runs +data +outputs diff --git a/dodgeballtracker/data_dumper.py b/dodgeballtracker/data_dumper.py new file mode 100644 index 0000000..bbdba84 --- /dev/null +++ b/dodgeballtracker/data_dumper.py @@ -0,0 +1,100 @@ +import abc + +import pandas as pd +from pathlib import Path +from dodgeballtracker.typing import TrackedDetectedObject, VideoTimestamp +from sqlite3 import Connection + +import pandera as pa + + +class DataDumper(abc.ABC): + + @abc.abstractmethod + def write_boxes(self, obj: list[TrackedDetectedObject], ts: VideoTimestamp, + frame_index: int) -> None: + pass + + @abc.abstractmethod + def start(self, video_path: Path, hexa_hash: str) -> None: + pass + + +class SQLiteDataDumper(DataDumper): + def __init__(self, conn: Connection): + self._conn = conn + self._video_id = None + + def _create_tables(self): + cur = self._conn.cursor() + + cur.executescript(''' + CREATE TABLE IF NOT EXISTS video ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + file_name TEXT NOT NULL, + file_sha256_hex TEXT NOT NULL, + + UNIQUE(file_sha256_hex) + ); + ''') + + cur.executescript(''' + CREATE TABLE IF NOT EXISTS video_frame( + id INTEGER PRIMARY KEY AUTOINCREMENT, + timestamp_ DOUBLE PRECISION NOT NULL, + index_ INTEGER NOT NULL, + video_id INTEGER NOT NULL, + + FOREIGN KEY(video_id) REFERENCES video(id), + UNIQUE(video_id, index_) + ); + ''') + + cur.execute(''' + CREATE TABLE IF NOT EXISTS tracked_detected_object( + id INTEGER PRIMARY KEY AUTOINCREMENT, + xmin DOUBLE PRECISION NOT NULL, + ymin DOUBLE PRECISION NOT NULL, + xmax DOUBLE PRECISION NOT NULL, + ymax DOUBLE PRECISION NOT NULL, + score DOUBLE PRECISION NOT NULL, + category INTEGER NOT NULL, + video_frame_id INTEGER NOT NULL, + + FOREIGN KEY(video_frame_id) REFERENCES video_frame(id) + ); + ''') + + self._conn.commit() + + def start(self, video_path: Path, hexa_hash: str) -> None: + self._create_tables() + cur = self._conn.cursor() + res = cur.execute(""" + INSERT INTO video(file_name, file_sha256_hex) + VALUES (?, ?) + ON CONFLICT(file_sha256_hex) DO NOTHING + RETURNING id + """, (video_path.name, hexa_hash)) + self._video_id = res.fetchone()[0] + self._conn.commit() + + @pa.check_io(df=TrackedDetectedObject) + def write_boxes(self, + df: pd.DataFrame, + ts: VideoTimestamp, + frame_index: int) -> None: + cur = self._conn.cursor() + res = cur.execute(''' + INSERT INTO video_frame(timestamp_, index_, video_id ) + VALUES (?, ?, ?) + ON CONFLICT(video_id, index_) DO NOTHING + RETURNING id + ''', (ts, frame_index, self._video_id)) + frame_id = res.fetchone()[0] + cols = ['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category'] + cur.executemany(''' + INSERT INTO tracked_detected_object(xmin, ymin, xmax, ymax, score, category, video_frame_id) + VALUES (?, ?, ?, ?, ?, ?, ?) + ''', (list(x[cols]) + [frame_id] for x in df.iloc)) + self._conn.commit() diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 4d4a746..3375749 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -10,7 +10,7 @@ import numpy as np import pandas as pd import pandera as pa from boxmot import DeepOcSort as DeepOCSORT -from .typing import VideoFrame, DetectedObject_, TrackedDetectedObject +from .typing import VideoFrame, DetectedObject, TrackedDetectedObject logger = logging.getLogger(__name__) @@ -107,7 +107,7 @@ class DodgeballTracker: return True return False - @pa.check_output(DetectedObject_) + @pa.check_output(DetectedObject) def run_model_on_frame(self, model, frame: VideoFrame): model_outputs = model(frame)[0] boxes = model_outputs.boxes.data.cpu().numpy() @@ -134,7 +134,7 @@ class DodgeballTracker: '/home/didier/all-dressed-programming/dodgeball-tracker/dodgeballtracker/runs/detect/train31/weights/best.pt') return YOLO(my_model, verbose=False) - @pa.check_io(detected_objs=DetectedObject_, out=TrackedDetectedObject) + @pa.check_io(detected_objs=DetectedObject, out=TrackedDetectedObject) def update_tracker(self, tracker: DeepOCSORT, detected_objs: pd.DataFrame, frame: VideoFrame): data = tracker.update(detected_objs.to_numpy(), frame) columns = ['xmin', 'ymin', 'xmax', 'ymax', 'object_id', 'score', 'category', 'index'] diff --git a/dodgeballtracker/test_data_dumper.py b/dodgeballtracker/test_data_dumper.py new file mode 100644 index 0000000..3422248 --- /dev/null +++ b/dodgeballtracker/test_data_dumper.py @@ -0,0 +1,82 @@ +import sqlite3 +import tempfile + +from pathlib import Path +from sqlite3 import Connection +from dodgeballtracker.typing import TrackedDetectedObject + +import pandas as pd +import pandera as pa +import numpy as np +import pytest +from contextlib import closing, contextmanager + +from dodgeballtracker.data_dumper import SQLiteDataDumper + + +def list_tables(conn: Connection) -> set[str]: + cur = conn.cursor() + res = cur.execute(''' + SELECT + name + FROM + sqlite_schema + WHERE + type ='table' AND + name NOT LIKE 'sqlite_%'; + ''') + return {row[0] for row in res} + +@pytest.fixture +def tmp_dir(): + with tempfile.TemporaryDirectory() as dir_: + yield Path(dir_) + + +def number_of_items_in_table(table: str, conn: Connection) -> int: + cur = conn.cursor() + res = cur.execute(f'SELECT COUNT(*) FROM {table}') + return res.fetchone()[0] + + +# @pytest.fixture +@contextmanager +def make_conn(tmp_dir: Path): + db_file= tmp_dir / 'db.sqlite' + conn_ = sqlite3.connect(db_file) + with closing(conn_) as open_conn: + open_conn.execute('PRAGMA foreign_keys = 1') + yield open_conn + + +@pytest.fixture +@pa.check_output(TrackedDetectedObject) +def mock_tracked_detected_object(): + rd = np.random.RandomState(seed=109) + data = rd.rand(10, len(TrackedDetectedObject.columns)) + return pd.DataFrame(data, columns=TrackedDetectedObject.columns) + + +def test_start(tmp_dir: Path): + with make_conn(tmp_dir) as conn: + dumper = SQLiteDataDumper(conn) + dumper.start(Path('foo/bar.MOV'), 'aaabbb') + assert dumper._video_id == 1 + + with make_conn(tmp_dir) as conn: + expt = { + 'tracked_detected_object', + 'video', + 'video_frame' + } + assert expt == list_tables(conn) + +# +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') + dumper.write_boxes(mock_tracked_detected_object, 40018, 34) + + with make_conn(tmp_dir) as conn: + assert 10 == number_of_items_in_table('tracked_detected_object', conn) \ No newline at end of file diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index d3262da..0ab2fc7 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -10,11 +10,11 @@ def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema: return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols}) -DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score", +DetectedObject = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score", "category"]) TrackedDetectedObject = tensor_of(Float64, - list(DetectedObject_.columns) + ["object_id", "index"]) + list(DetectedObject.columns) + ["object_id", "index"]) Color = tuple[int, int, int] diff --git a/flake.nix b/flake.nix index 5d92b2a..9497d1f 100644 --- a/flake.nix +++ b/flake.nix @@ -62,6 +62,7 @@ ps.matplotlib ps.torchWithCuda ps.jupyter + ps.pytest pythonPackages.opencv4 pythonPackages.boxmot pythonPackages.ultralytics