all-dressed-programming/dodgeball-tracker
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Going somewhere before going to the gym

Browse Source
This commit is contained in:
didier committed 2024-11-15 13:33:14 -05:00
1 parent f5b8e24e01
commit 236dd82871
4 files changed
+96 -56
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

dodgeballtracker/__main__.py
+31 -10
View File
@@ -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,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):
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__":
dodgeballtracker/data_dumper.py
+13 -4
View File
@@ -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]
dodgeballtracker/logic.py
+4 -41
View File
@@ -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
+48 -1
View File
@@ -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)