From d7017c756bfb7b466acf8ab329c226c33862fcd4 Mon Sep 17 00:00:00 2001 From: didier amyot Date: Sat, 2 Nov 2024 14:11:18 -0400 Subject: [PATCH] definitively going somewhere. --- .envrc | 1 + .gitignore | 1 + dodgeballtracker/__main__.py | 57 +++++++++++++++++++++++---- dodgeballtracker/helper.py | 17 ++------ dodgeballtracker/logic.py | 30 +++++++------- dodgeballtracker/video_writer.py | 67 ++++++++++++++++++++++++++++++++ flake.nix | 7 ++-- 7 files changed, 139 insertions(+), 41 deletions(-) create mode 100644 .envrc create mode 100644 dodgeballtracker/video_writer.py diff --git a/.envrc b/.envrc new file mode 100644 index 0000000..3550a30 --- /dev/null +++ b/.envrc @@ -0,0 +1 @@ +use flake diff --git a/.gitignore b/.gitignore index 6c14865..b46071c 100644 --- a/.gitignore +++ b/.gitignore @@ -159,3 +159,4 @@ env *.pt *.png *.ipynb +.direnv diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index 3a5a70d..a419fb2 100644 --- a/dodgeballtracker/__main__.py +++ b/dodgeballtracker/__main__.py @@ -1,28 +1,69 @@ +import time +from contextlib import closing from pathlib import Path +from dodgeballtracker.helper import time_tracker from dodgeballtracker.logic import (DetectedObjectBox, DodgeballTracker, get_court_perspective_transform) import numpy as np import logging +import click +from typing import ContextManager +from dodgeballtracker.video_writer import (MP4VVideoWriter, VideoWriter, + MP4VVideoWriterWithFrame) here = Path(__file__).parent input_dir = here / '..' / 'data' output_dir = here / '..' / 'outputs' -if __name__ == "__main__": + +@click.group() +@click.option('--debug/--no-debug', default=False) +def cli(debug): logging.basicConfig(filename='tracker.log', level=logging.DEBUG) - detected_box_objects = [ - DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05) - ] - input_video = input_dir / "small.mp4" - output_video = output_dir / "foo.mp4" +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]: + if display_frame: + return closing( + MP4VVideoWriterWithFrame.from_frame(video_input, video_output)) + return closing(MP4VVideoWriter.from_frame(video_input, video_output)) + + +@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() - dodgeball_tracker = DodgeballTracker(input_video, output_video, detected_box_objects, transformation_matrix) + with get_video_writer(display_video, input_, output) as video_writer: + tracker = DodgeballTracker(input_, + output, + video_writer, + detected_box_objects, + transformation_matrix) - dodgeball_tracker.main() + with time_tracker(): + tracker.main() + + +if __name__ == "__main__": + cli() diff --git a/dodgeballtracker/helper.py b/dodgeballtracker/helper.py index 924ae95..41b0747 100644 --- a/dodgeballtracker/helper.py +++ b/dodgeballtracker/helper.py @@ -1,21 +1,10 @@ import contextlib import time import cv2 +import logging +logger = logging.getLogger(__name__) -def create_video_writer(video_cap, output_filename): - - # grab the width, height, and fps of the frames in the video stream. - frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) - frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - fps = int(video_cap.get(cv2.CAP_PROP_FPS)) - - # initialize the FourCC and a video writer object - fourcc = cv2.VideoWriter_fourcc(*'MP4V') - writer = cv2.VideoWriter(output_filename, fourcc, fps, - (frame_width, frame_height)) - - return writer def iter_frames(video_cap): @@ -31,4 +20,4 @@ def time_tracker(): start = time.monotonic() yield end = time.monotonic() - print(f"Time to process 1 frame: {(end-start) * 1000:.0f} milliseconds") \ No newline at end of file + logger.debug(f"Time to process 1 frame: {(end-start) * 1000:.0f} milliseconds") diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 2477fae..d038dd2 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -4,7 +4,8 @@ import logging import torch from ultralytics import YOLO import cv2 -from dodgeballtracker.helper import create_video_writer, iter_frames +from dodgeballtracker.helper import iter_frames +from dodgeballtracker.video_writer import create_video_writer, VideoWriter import matplotlib import numpy as np from boxmot import DeepOcSort as DeepOCSORT @@ -15,15 +16,15 @@ logger = logging.getLogger(__name__) Color = tuple[int, int, int] -def get_best_available_torch_device(): +def get_best_available_torch_device() -> torch.device: if torch.cuda.is_available(): logger.debug("Using the cuda (gpu) torch device.") - return 'cuda' + return torch.device('cuda:0') if torch.backends.mps.is_available(): logger.debug("Using the mps (apple silicon) torch device.") - return 'mps' + return torch.device('mps') logger.debug("Default to cpu (no apple silicon or cuda) torch device.") - return 'cpu' + return torch.device('cpu') class DetectedObjectBox: @@ -66,16 +67,21 @@ def get_court_perspective_transform(): return cv2.getPerspectiveTransform(empirical_source, empirical_target) + + + class DodgeballTracker: def __init__(self, input_file: Path, output_file: Path, + video_writer: VideoWriter, 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.detected_box_objects = sorted(detected_box_objects, reverse=True) self.transformation_matrix = transformation_matrix self.points = [] @@ -138,7 +144,7 @@ class DodgeballTracker: def _get_tracker(self): return DeepOCSORT( Path('lmbn_n_cuhk03_d.pt'), # which ReID model to use - torch.device(get_best_available_torch_device()), + get_best_available_torch_device(), # torch.cuda.get_device_name(0), #torch.device(get_best_available_torch_device()), True, ) @@ -150,10 +156,8 @@ class DodgeballTracker: def main(self): # define some constants - print("started") - video_cap = cv2.VideoCapture(str(self.input_path)) - writer = create_video_writer(video_cap, str(self.output_path)) + # writer = create_video_writer(video_cap, self.output_path) tracker = self._get_tracker() # load the pre-trained YOLOv8n model @@ -168,19 +172,13 @@ class DodgeballTracker: if is_good_box: res.append(data) - # show the frame to our screen - print("tracker?") data_ = tracker.update(np.array(res), frame) self.print_tracker_data(data_, i, ts) - self.write_frame(writer, frame) + self.video_writer.write_frame(frame) if cv2.waitKey(1) == ord("q"): break - print("fini") video_cap.release() - writer.release() - cv2.destroyAllWindows() self.persist_points() - print("finished") diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py new file mode 100644 index 0000000..2d3468e --- /dev/null +++ b/dodgeballtracker/video_writer.py @@ -0,0 +1,67 @@ +import abc +from pathlib import Path + +import cv2 + +def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter: + + # grab the width, height, and fps of the frames in the video stream. + frame_width = int(video_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + frame_height = int(video_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = int(video_cap.get(cv2.CAP_PROP_FPS)) + + # initialize the FourCC and a video writer object + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + writer = cv2.VideoWriter(str(output_filename.absolute()), + fourcc, + fps, + (frame_width, frame_height)) + + return writer + + +class VideoWriter(abc.ABC): + + @abc.abstractmethod + def write_frame(self, frame) -> None: + pass + + @abc.abstractmethod + def close(self) -> None: + pass + + +class NoopVideoWriter(VideoWriter): + + def write_frame(self, frame) -> None: + pass + + def close(self) -> None: + pass + + +class MP4VVideoWriter(VideoWriter): + + def __init__(self, video_writer: cv2.VideoWriter) -> None: + self._video_writer = video_writer + + def write_frame(self, frame) -> None: + self._video_writer.write(frame) + + def close(self) -> None: + self._video_writer.release() + cv2.destroyAllWindows() + + @classmethod + def from_frame(cls, input_: Path, output: Path) -> VideoWriter: + video_cap = cv2.VideoCapture(str(input_)) + try: + return cls(create_video_writer(video_cap, output)) + finally: + video_cap.release() + + +class MP4VVideoWriterWithFrame(MP4VVideoWriter): + def write_frame(self, frame) -> None: + super().write_frame(frame) + cv2.imshow("Frame", frame) diff --git a/flake.nix b/flake.nix index 4399f74..b7e2372 100644 --- a/flake.nix +++ b/flake.nix @@ -27,9 +27,10 @@ virtual-env = pkgs.python312.withPackages (ps: [ ps.matplotlib ps.torchWithCuda - # pythonPackages.opencv4 - # pythonPackages.boxmot - # pythonPackages.ultralytics + pythonPackages.opencv4 + pythonPackages.boxmot + pythonPackages.ultralytics + pythonPackages.click ]); in {