Going somewhere before going to the gym

dodgeballtracker/__main__.py

@@ -1,12 +1,15 @@
-from contextlib import closing
+from contextlib import closing, contextmanager
 from pathlib import Path
+import sqlite3
 
 from dodgeballtracker.helper import time_tracker
 from dodgeballtracker.logic import (DodgeballTracker,
                                     get_court_perspective_transform)
+from dodgeballtracker.data_dumper import SQLiteDataDumper
 import numpy as np
 import logging
 import click
+import hashlib
 from typing import ContextManager
 
 from dodgeballtracker.video_writer import (MP4VVideoWriter,
@@ -25,10 +28,13 @@ def cli(debug):
     logging.basicConfig(filename='tracker.log', level=logging.DEBUG)
 
 
-def get_output(input_: Path, proposed_output: Path | None) -> Path:
+def get_output_dir(input_: Path, proposed_output: Path | None) -> Path:
     if proposed_output is not None:
+        proposed_output.mkdir(parents=True, exist_ok=True)
         return proposed_output
-    return output_dir / input_.name
+    output = output_dir / input_.name
+    output.mkdir(parents=True, exist_ok=True)
+    return output
 
 
 def get_video_writer(display_frame: bool,
@@ -41,8 +47,22 @@ def get_video_writer(display_frame: bool,
     return closing(MP4VVideoWriter.from_frame(video_input, video_output, thresholds))
 
 
-def process_logic(input_: Path, output: Path | None, display_video: bool):
+def get_file_sha256(file_: Path) -> str:
-    output = get_output(input_, output)
+    with open(file_, 'rb', buffering=0) as f:
+        return hashlib.file_digest(f, 'sha256').hexdigest()
+
+
+@contextmanager
+def create_data_dumper(input_file: Path, outout_file: Path):
+    with closing(sqlite3.connect(outout_file)) as conn:
+        conn.execute('PRAGMA foreign_keys = 1')
+        dumper = SQLiteDataDumper(conn)
+        dumper.start(input_file, get_file_sha256(input_file))
+        yield dumper
+
+
+def process_logic(input_: Path, output_dir: Path | None, display_video: bool):
+    output_dir = get_output_dir(input_, output_dir)
     detected_box_objects = [
         DetectedObjectBox.from_viridis(threshold) for threshold in
         np.arange(0.5, 1, 0.05)
@@ -50,10 +70,11 @@ def process_logic(input_: Path, output: Path | None, display_video: bool):
 
     transformation_matrix = get_court_perspective_transform()
 
-    with get_video_writer(display_video, input_, output) as video_writer:
+    with get_video_writer(display_video, input_, output_dir / 'video.mov') as video_writer:
+        with create_data_dumper(input_, output_dir / 'data.sqlite') as data_dumper:
             tracker = DodgeballTracker(input_,
-                                   output,
                                        video_writer,
+                                       data_dumper,
                                        detected_box_objects,
                                        transformation_matrix)
 
@@ -68,10 +89,10 @@ def train():
 @cli.command()
 @click.option('--input', '-i', 'input_',
               type=click.Path(readable=True, path_type=Path))
-@click.option('--output', '-o', type=click.Path(path_type=Path))
+@click.option('--output-dir', '-o', type=click.Path(path_type=Path))
 @click.option('--display-video', '-d', type=bool,  is_flag=True, show_default=True, default=False)
-def process(input_: Path, output: Path | None, display_video: bool):
+def process(input_: Path, output_dir: Path | None, display_video: bool):
-    process_logic(input_, output, display_video)
+    process_logic(input_, output_dir, display_video)
 
 
 if __name__ == "__main__":

dodgeballtracker/data_dumper.py

@@ -33,6 +33,7 @@ class SQLiteDataDumper(DataDumper):
                     id INTEGER PRIMARY KEY AUTOINCREMENT,
                     file_name TEXT NOT NULL,
                     file_sha256_hex TEXT NOT NULL,
+                    count INTEGER NOT NULL DEFAULT 1,
                     
                     UNIQUE(file_sha256_hex)
                 );
@@ -44,6 +45,7 @@ class SQLiteDataDumper(DataDumper):
                                     timestamp_ DOUBLE PRECISION NOT NULL,
                                     index_ INTEGER NOT NULL,
                                     video_id INTEGER NOT NULL,
+                                    nbr_detected_object INTEGER NOT NULL,
                                     
                                     FOREIGN KEY(video_id) REFERENCES video(id),
                                     UNIQUE(video_id, index_)
@@ -69,26 +71,33 @@ class SQLiteDataDumper(DataDumper):
 
     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
+        ON CONFLICT(file_sha256_hex) DO UPDATE SET count = count + 1
         RETURNING id
         """, (video_path.name, hexa_hash))
         self._video_id = res.fetchone()[0]
         self._conn.commit()
 
+    def _is_empty_df(self, df):
+        return not df.shape[0]
+
     @pa.check_io(df=TrackedDetectedObject)
     def write_boxes(self,
                     df: pd.DataFrame,
                     ts: VideoTimestamp,
                     frame_index: int) -> None:
+        if self._is_empty_df(df):
+            return
         cur = self._conn.cursor()
         res = cur.execute('''
-        INSERT INTO video_frame(timestamp_, index_, video_id )
+        INSERT INTO video_frame(timestamp_, index_, video_id, nbr_detected_object )
-        VALUES (?, ?, ?)
+        VALUES (?, ?, ?, 1)
-        ON CONFLICT(video_id, index_) DO NOTHING
+            ON CONFLICT(video_id, index_) 
+                DO UPDATE SET nbr_detected_object = nbr_detected_object + 1
         RETURNING id
         ''', (ts, frame_index, self._video_id))
         frame_id = res.fetchone()[0]

dodgeballtracker/logic.py

@@ -6,6 +6,7 @@ from ultralytics import YOLO
 import cv2
 from dodgeballtracker.helper import iter_frames
 from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox
+from dodgeballtracker.data_dumper import DataDumper
 import numpy as np
 import pandas as pd
 import pandera as pa
@@ -60,52 +61,17 @@ class DodgeballTracker:
 
     def __init__(self,
                  input_file: Path,
-                 output_file: Path,
                  video_writer: VideoWriter,
+                 data_dumper: DataDumper,
                  detected_box_objects: list[DetectedObjectBox],
                  transformation_matrix):
         self.input_path = input_file
-        self.output_path = output_file
-        self.output_json_path = output_file.with_suffix(".json")
         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.points = []
 
-    def write_frame(self, writer, frame) -> None:
-        cv2.imshow("Frame", frame)
-        writer.write(frame)
-
-    def extract_box(self, data) -> tuple[int, int, int, int]:
-        return int(data[0]), int(data[1]), int(data[2]), int(data[3])
-
-    def print_box(self, data, frame_idx, frame_ts):
-        xmin, ymin, xmax, ymax = self.extract_box(data)
-        xavg = (xmax + xmin) / 2
-        v = np.array([xavg, ymax, 1])
-        ret = self.transformation_matrix @ v
-        x, y = ret[0] / ret[2], ret[1] / ret[2]
-        object_id = data[-1]
-        self.points.append({"x": x,
-                            "y": y,
-                            "frameIndex": frame_idx,
-                            "frameTimestamp": frame_ts,
-                            "objectId": int(object_id)
-                            })
-
-    def print_tracker_data(self, tracker_data, frame_idx, frame_ts):
-        for d in tracker_data:
-            self.print_box(d, frame_idx, frame_ts)
-
-    def plot_box_to_frame(self, data, frame: VideoFrame):
-        confidence = float(data[4])
-        for detected_box_object in self.detected_box_objects:
-
-            if detected_box_object.min_threshold < confidence:
-                # xmin, ymin, xmax, ymax = self.extract_box(data)
-                # cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2)
-                return True
-        return False
 
     @pa.check_output(DetectedObject)
     def run_model_on_frame(self, model, frame: VideoFrame):
@@ -115,9 +81,6 @@ class DodgeballTracker:
         ret = pd.DataFrame(boxes, columns=columns)
         return ret
 
-    def persist_points(self):
-        with open(self.output_json_path, "wt") as file_:
-            json.dump({"points": self.points}, file_, indent=4)
 
     def _get_tracker(self) -> DeepOCSORT:
         return DeepOCSORT(
@@ -155,6 +118,7 @@ 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.video_writer.plot_boxes_to_frame(tracked_obj, frame)
             self.video_writer.write_frame(frame)
 
@@ -162,4 +126,3 @@ class DodgeballTracker:
                 break
 
         video_cap.release()
-        self.persist_points()

dodgeballtracker/test_data_dumper.py

@@ -71,7 +71,25 @@ def test_start(tmp_dir: Path):
         }
         assert expt == list_tables(conn)
 
-#
+
+def test_restart(tmp_dir: Path):
+    with make_conn(tmp_dir) as conn:
+        dumper = SQLiteDataDumper(conn)
+        dumper.start(Path('foo/bar.MOV'), 'aaabbb')
+
+    with make_conn(tmp_dir) as conn:
+        dumper = SQLiteDataDumper(conn)
+        dumper.start(Path('foo/bar.MOV'), 'aaabbb')
+
+    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)
@@ -80,3 +98,32 @@ def test_write_boxes_simple(tmp_dir: Path, mock_tracked_detected_object: pd.Data
 
     with make_conn(tmp_dir) as conn:
         assert 10 == number_of_items_in_table('tracked_detected_object', conn)
+
+
+def test_write_boxes_multiple_times(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:
+        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 20 == number_of_items_in_table('tracked_detected_object', conn)
+
+
+
+def test_write_boxes_with_empty_frame(tmp_dir: Path,
+                                      mock_tracked_detected_object: pd.DataFrame):
+    empty_tracked_detected_object = mock_tracked_detected_object[0:0]
+    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)
+
+    with make_conn(tmp_dir) as conn:
+        assert 0 == number_of_items_in_table('tracked_detected_object', conn)
+        assert 0 == number_of_items_in_table('video_frame', conn)