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