125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
import json5 as json
|
|
from contextlib import closing, contextmanager
|
|
from pathlib import Path
|
|
import sqlite3
|
|
from dodgeballtracker.metadata import VideoMetadata
|
|
|
|
|
|
from dodgeballtracker.helper import time_tracker
|
|
from dodgeballtracker.logic import (DodgeballTracker,
|
|
get_court_perspective_transform)
|
|
from dodgeballtracker.data_dumper import SQLiteDataDumper, DataDumper
|
|
import numpy as np
|
|
import logging
|
|
import click
|
|
import hashlib
|
|
from typing import ContextManager, Iterator
|
|
|
|
from dodgeballtracker.video_writer import (MP4VVideoWriter,
|
|
VideoWriter,
|
|
DetectedObjectBox,
|
|
MP4VVideoWriterWithFrame)
|
|
|
|
|
|
|
|
@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,
|
|
output_file: Path,
|
|
metadata: VideoMetadata) -> Iterator[DataDumper]:
|
|
with closing(sqlite3.connect(output_file)) as conn:
|
|
conn.execute('PRAGMA foreign_keys = 1')
|
|
dumper = SQLiteDataDumper(conn)
|
|
dumper.start(input_file, get_file_sha256(input_file), metadata)
|
|
yield dumper
|
|
|
|
def deserialized_video_metadata(positional_data: Path | None) -> VideoMetadata | None:
|
|
if not positional_data:
|
|
return None
|
|
with open(positional_data, 'rt') as file_:
|
|
data = json.load(file_)
|
|
return VideoMetadata(**data)
|
|
|
|
|
|
def process_logic(input_: Path,
|
|
output_dir: Path | None,
|
|
metadata_file: Path | None,
|
|
display_video: bool):
|
|
metadata = deserialized_video_metadata(metadata_file)
|
|
|
|
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', metadata) 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('--metadata-file', '-f', type=click.Path(path_type=Path, file_okay=True, exists=True))
|
|
@click.option('--display-video', '-d', type=bool, is_flag=True, show_default=True, default=False)
|
|
def process(input_: Path,
|
|
output_dir: Path | None,
|
|
metadata_file: Path | None,
|
|
display_video: bool):
|
|
process_logic(input_,
|
|
output_dir,
|
|
metadata_file,
|
|
display_video)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|
|
# process_logic(Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/experiments/rue-honore-beaugrand-ouest-001.MOV'),
|
|
# Path('/home/didier/all-dressed-programming/dodgeball-tracker/outputs/benchmark-001'),
|
|
# False)
|