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):
-    output = get_output(input_, output)
+def get_file_sha256(file_: Path) -> str:
+    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,15 +70,16 @@ 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:
-        tracker = DodgeballTracker(input_,
-                                   output,
-                                   video_writer,
-                                   detected_box_objects,
-                                   transformation_matrix)
+    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_,
+                                       video_writer,
+                                       data_dumper,
+                                       detected_box_objects,
+                                       transformation_matrix)
 
-        with time_tracker():
-            tracker.main()
+            with time_tracker():
+                tracker.main()
 
 
 @cli.command()
@@ -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):
-    process_logic(input_, output, display_video)
+def process(input_: Path, output_dir: Path | None, display_video: bool):
+    process_logic(input_, output_dir, display_video)
 
 
 if __name__ == "__main__":