From 152b3b9da534b35c96c5c59315781254a3f7cbde Mon Sep 17 00:00:00 2001 From: didier amyot Date: Sat, 2 Nov 2024 14:57:59 -0400 Subject: [PATCH] definitively going somewhere. --- dodgeballtracker/helper.py | 4 +++- dodgeballtracker/logic.py | 29 ++++++++--------------------- dodgeballtracker/typing.py | 7 +++++++ dodgeballtracker/video_writer.py | 23 ++++++++++++++++++----- 4 files changed, 36 insertions(+), 27 deletions(-) create mode 100644 dodgeballtracker/typing.py diff --git a/dodgeballtracker/helper.py b/dodgeballtracker/helper.py index 41b0747..4fac932 100644 --- a/dodgeballtracker/helper.py +++ b/dodgeballtracker/helper.py @@ -1,5 +1,7 @@ import contextlib import time +from typing import Iterator +from dodgeballtracker.typing import VideoFrame, VideoTimestamp import cv2 import logging @@ -7,7 +9,7 @@ logger = logging.getLogger(__name__) -def iter_frames(video_cap): +def iter_frames(video_cap) -> Iterator[tuple[VideoFrame, VideoTimestamp]]: while 1: ret, frame = video_cap.read() if not ret: diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index d038dd2..bc021f8 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -1,4 +1,5 @@ import json +from collections import Counter from pathlib import Path import logging import torch @@ -9,7 +10,7 @@ from dodgeballtracker.video_writer import create_video_writer, VideoWriter import matplotlib import numpy as np from boxmot import DeepOcSort as DeepOCSORT - +from .typing import VideoFrame logger = logging.getLogger(__name__) @@ -111,7 +112,7 @@ class DodgeballTracker: for d in tracker_data: self.print_box(d, frame_idx, frame_ts) - def plot_box_to_frame(self, data, frame): + def plot_box_to_frame(self, data, frame: VideoFrame): confidence = float(data[4]) for detected_box_object in self.detected_box_objects: @@ -121,22 +122,9 @@ class DodgeballTracker: return True return False - def run_model_on_frame(self, model, frame): + def run_model_on_frame(self, model, frame: VideoFrame): return model(frame)[0] - def crazy_experiment(self, frame): - cv2.putText(frame, - "Hello world", - (50, 50), - cv2.FONT_HERSHEY_SIMPLEX, - 2, - (0, 255, 0), 8) - cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5) - pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32) - pts = pts.reshape((-1, 1, 2)) - - cv2.polylines(frame, pts, True, (0, 255, 255), 8) - def persist_points(self): with open(self.output_json_path, "wt") as file_: json.dump({"points": self.points}, file_, indent=4) @@ -144,7 +132,7 @@ class DodgeballTracker: def _get_tracker(self): return DeepOCSORT( Path('lmbn_n_cuhk03_d.pt'), # which ReID model to use - get_best_available_torch_device(), # torch.cuda.get_device_name(0), #torch.device(get_best_available_torch_device()), + get_best_available_torch_device(), True, ) @@ -157,11 +145,9 @@ class DodgeballTracker: def main(self): # define some constants video_cap = cv2.VideoCapture(str(self.input_path)) - # writer = create_video_writer(video_cap, self.output_path) tracker = self._get_tracker() - - # load the pre-trained YOLOv8n model model = self._get_detection_model() + counter = Counter() for i, (frame, ts) in enumerate(iter_frames(video_cap)): # run the YOLO model on the frame @@ -177,8 +163,9 @@ class DodgeballTracker: self.video_writer.write_frame(frame) - if cv2.waitKey(1) == ord("q"): + if self.video_writer.wait(): break + print(f"counter: {counter}") video_cap.release() self.persist_points() diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py new file mode 100644 index 0000000..0a011ea --- /dev/null +++ b/dodgeballtracker/typing.py @@ -0,0 +1,7 @@ +from typing import NewType + +type XCoord = float +type YCoord = NewType('YCoord', float) +type Point = tuple[XCoord, YCoord] +type VideoFrame = NewType('Frame', float) +type VideoTimestamp = float diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py index 2d3468e..c699a56 100644 --- a/dodgeballtracker/video_writer.py +++ b/dodgeballtracker/video_writer.py @@ -1,6 +1,6 @@ import abc from pathlib import Path - +from dodgeballtracker.typing import VideoFrame import cv2 def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> cv2.VideoWriter: @@ -23,35 +23,45 @@ def create_video_writer(video_cap: cv2.VideoCapture, output_filename: Path) -> c class VideoWriter(abc.ABC): @abc.abstractmethod - def write_frame(self, frame) -> None: + def write_frame(self, frame: VideoFrame) -> None: pass @abc.abstractmethod def close(self) -> None: pass + @abc.abstractmethod + def wait(self) -> bool: + pass + class NoopVideoWriter(VideoWriter): - def write_frame(self, frame) -> None: + def write_frame(self, frame: VideoFrame) -> None: pass def close(self) -> None: pass + def wait(self) -> bool: + return False + class MP4VVideoWriter(VideoWriter): def __init__(self, video_writer: cv2.VideoWriter) -> None: self._video_writer = video_writer - def write_frame(self, frame) -> None: + def write_frame(self, frame: VideoFrame) -> None: self._video_writer.write(frame) def close(self) -> None: self._video_writer.release() cv2.destroyAllWindows() + def wait(self) -> bool: + return False #cv2.waitKey(1) == ord("q") + @classmethod def from_frame(cls, input_: Path, output: Path) -> VideoWriter: video_cap = cv2.VideoCapture(str(input_)) @@ -62,6 +72,9 @@ class MP4VVideoWriter(VideoWriter): class MP4VVideoWriterWithFrame(MP4VVideoWriter): - def write_frame(self, frame) -> None: + def write_frame(self, frame: VideoFrame) -> None: super().write_frame(frame) cv2.imshow("Frame", frame) + + def wait(self) -> bool: + return cv2.waitKey(1) == ord("q")