From 236dd828713d3901fd1bde673dadeeffab76c484 Mon Sep 17 00:00:00 2001 From: didier amyot Date: Fri, 15 Nov 2024 13:33:14 -0500 Subject: [PATCH] Going somewhere before going to the gym --- dodgeballtracker/__main__.py | 53 +++++++++++++++++++--------- dodgeballtracker/data_dumper.py | 17 ++++++--- dodgeballtracker/logic.py | 45 +++-------------------- dodgeballtracker/test_data_dumper.py | 51 ++++++++++++++++++++++++-- 4 files changed, 103 insertions(+), 63 deletions(-) diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index 2cc0a83..02c86bf 100644 --- a/dodgeballtracker/__main__.py +++ b/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__": diff --git a/dodgeballtracker/data_dumper.py b/dodgeballtracker/data_dumper.py index bbdba84..8e90b95 100644 --- a/dodgeballtracker/data_dumper.py +++ b/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 ) - VALUES (?, ?, ?) - ON CONFLICT(video_id, index_) DO NOTHING + INSERT INTO video_frame(timestamp_, index_, video_id, nbr_detected_object ) + VALUES (?, ?, ?, 1) + 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] diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 3375749..7a9d5e8 100644 --- a/dodgeballtracker/logic.py +++ b/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() diff --git a/dodgeballtracker/test_data_dumper.py b/dodgeballtracker/test_data_dumper.py index 3422248..e46222b 100644 --- a/dodgeballtracker/test_data_dumper.py +++ b/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) @@ -79,4 +97,33 @@ def test_write_boxes_simple(tmp_dir: Path, mock_tracked_detected_object: pd.Data 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 + 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) \ No newline at end of file