101 lines
3.5 KiB
Python
101 lines
3.5 KiB
Python
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,
|
|
VideoWriter,
|
|
DetectedObjectBox,
|
|
MP4VVideoWriterWithFrame)
|
|
|
|
here = Path(__file__).parent
|
|
input_dir = here / '..' / 'data'
|
|
output_dir = here / '..' / 'outputs'
|
|
|
|
|
|
@click.group()
|
|
@click.option('--debug/--no-debug', default=False)
|
|
def cli(debug):
|
|
logging.basicConfig(filename='tracker.log', level=logging.DEBUG)
|
|
|
|
|
|
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
|
|
output = output_dir / input_.name
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
return output
|
|
|
|
|
|
def get_video_writer(display_frame: bool,
|
|
video_input: Path,
|
|
video_output: Path) -> ContextManager[VideoWriter]:
|
|
thresholds = list(np.arange(0.5, 1, 0.05))
|
|
if display_frame:
|
|
return closing(
|
|
MP4VVideoWriterWithFrame.from_frame(video_input, video_output, thresholds))
|
|
return closing(MP4VVideoWriter.from_frame(video_input, video_output, thresholds))
|
|
|
|
|
|
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)
|
|
]
|
|
|
|
transformation_matrix = get_court_perspective_transform()
|
|
|
|
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()
|
|
|
|
|
|
@cli.command()
|
|
def train():
|
|
pass
|
|
|
|
@cli.command()
|
|
@click.option('--input', '-i', 'input_',
|
|
type=click.Path(readable=True, 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_dir: Path | None, display_video: bool):
|
|
process_logic(input_, output_dir, display_video)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|
|
# process_logic(Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/experiments/rue-honore-beaugrand-ouest-001.MOV'), None, False)
|