Going somewhere with a coffee.

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

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']

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)

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]