71 lines
2.4 KiB
Python
71 lines
2.4 KiB
Python
from contextlib import closing
|
|
from pathlib import Path
|
|
|
|
from dodgeballtracker.helper import time_tracker
|
|
from dodgeballtracker.logic import (DodgeballTracker,
|
|
get_court_perspective_transform)
|
|
import numpy as np
|
|
import logging
|
|
import click
|
|
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(input_: Path, proposed_output: Path | None) -> Path:
|
|
if proposed_output is not None:
|
|
return proposed_output
|
|
return output_dir / input_.name
|
|
|
|
|
|
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))
|
|
|
|
|
|
@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('--display-video', '-d', type=bool, is_flag=True, show_default=True, default=False)
|
|
def process(input_: Path, output: Path | None, display_video: bool):
|
|
output = get_output(input_, output)
|
|
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) as video_writer:
|
|
tracker = DodgeballTracker(input_,
|
|
output,
|
|
video_writer,
|
|
detected_box_objects,
|
|
transformation_matrix)
|
|
|
|
with time_tracker():
|
|
tracker.main()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli()
|