diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index a419fb2..10762c2 100644 --- a/dodgeballtracker/__main__.py +++ b/dodgeballtracker/__main__.py @@ -1,17 +1,17 @@ -import time from contextlib import closing from pathlib import Path from dodgeballtracker.helper import time_tracker -from dodgeballtracker.logic import (DetectedObjectBox, - DodgeballTracker, +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, +from dodgeballtracker.video_writer import (MP4VVideoWriter, + VideoWriter, + DetectedObjectBox, MP4VVideoWriterWithFrame) here = Path(__file__).parent @@ -34,10 +34,11 @@ def get_output(input_: Path, proposed_output: Path | None) -> Path: 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)) - return closing(MP4VVideoWriter.from_frame(video_input, video_output)) + MP4VVideoWriterWithFrame.from_frame(video_input, video_output, thresholds)) + return closing(MP4VVideoWriter.from_frame(video_input, video_output, thresholds)) @cli.command() diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index bc021f8..ca5e11d 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -1,16 +1,14 @@ import json -from collections import Counter from pathlib import Path import logging import torch from ultralytics import YOLO import cv2 from dodgeballtracker.helper import iter_frames -from dodgeballtracker.video_writer import create_video_writer, VideoWriter -import matplotlib +from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox import numpy as np from boxmot import DeepOcSort as DeepOCSORT -from .typing import VideoFrame +from .typing import VideoFrame, DetectedObject logger = logging.getLogger(__name__) @@ -28,21 +26,6 @@ def get_best_available_torch_device() -> torch.device: return torch.device('cpu') -class DetectedObjectBox: - def __init__(self, min_threshold: float, color: Color): - self.min_threshold = min_threshold - self.color = color - - def __lt__(self, other): - return self.min_threshold < other.min_threshold - - @classmethod - def from_viridis(cls, min_threshold): - cm = matplotlib.colormaps['viridis'] - color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3))) - return cls(min_threshold, color) - - def get_court_perspective_transform(): """ """ @@ -117,13 +100,17 @@ class DodgeballTracker: for detected_box_object in self.detected_box_objects: if detected_box_object.min_threshold < confidence: - xmin, ymin, xmax, ymax = self.extract_box(data) - cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2) + # xmin, ymin, xmax, ymax = self.extract_box(data) + # cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2) return True return False def run_model_on_frame(self, model, frame: VideoFrame): - return model(frame)[0] + model_outputs = model(frame)[0] + for data in model_outputs.boxes.data.tolist(): + detected_obj = DetectedObject(*data) + if detected_obj.confidence > 0.5: + yield detected_obj def persist_points(self): with open(self.output_json_path, "wt") as file_: @@ -142,30 +129,28 @@ class DodgeballTracker: """ return YOLO("yolo11n.pt", verbose=True) + def update_tracker(self, tracker, detected_objs: list[DetectedObject], frame: VideoFrame) -> None: + important_detected_objs = [do for do in detected_objs if do.confidence > 0.5] + def main(self): # define some constants video_cap = cv2.VideoCapture(str(self.input_path)) tracker = self._get_tracker() model = self._get_detection_model() - counter = Counter() for i, (frame, ts) in enumerate(iter_frames(video_cap)): # run the YOLO model on the frame - res = [] - detections = self.run_model_on_frame(model, frame) - for data in detections.boxes.data.tolist(): - is_good_box = self.plot_box_to_frame(data, frame) - if is_good_box: - res.append(data) - - data_ = tracker.update(np.array(res), frame) - self.print_tracker_data(data_, i, ts) + detected_objs = list(self.run_model_on_frame(model, frame)) + self.video_writer.plot_boxes_to_frame(detected_objs, frame) + # data_ = tracker.update(np.array([dobj.as_nparray() for dobj in detected_objs]), frame) + # if data_: + # print("dkld") + # self.print_tracker_data(data_, i, ts) self.video_writer.write_frame(frame) 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 index 0a011ea..334802f 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -1,7 +1,37 @@ from typing import NewType +from dataclasses import dataclass +import numpy as np -type XCoord = float -type YCoord = NewType('YCoord', float) -type Point = tuple[XCoord, YCoord] -type VideoFrame = NewType('Frame', float) -type VideoTimestamp = float +XCoord = NewType('XCoord', float) +YCoord = NewType('YCoord', float) +Confidence = NewType('Confidence', float) +Category = NewType('Category', float) + + + +@dataclass +class Point: + x: XCoord + y: YCoord + +VideoFrame = NewType('Frame', float) +VideoTimestamp = float + +@dataclass +class DetectedObject: + xmin: XCoord + ymin: YCoord + xmax: XCoord + ymax: YCoord + confidence: Confidence + category: Category + + def as_int_box(self) -> tuple[tuple[int, int], tuple[int, int]]: + return (int(self.xmin), int(self.ymin)), (int(self.xmax), int(self.ymax)) + + def as_nparray(self): + return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.confidence, self.category]) + + + +Color = tuple[int, int, int] \ No newline at end of file diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py index c699a56..293e8f6 100644 --- a/dodgeballtracker/video_writer.py +++ b/dodgeballtracker/video_writer.py @@ -1,8 +1,25 @@ import abc from pathlib import Path -from dodgeballtracker.typing import VideoFrame +from dodgeballtracker.typing import VideoFrame, DetectedObject, Color +import matplotlib import cv2 + +class DetectedObjectBox: + def __init__(self, min_threshold: float, color: Color): + self.min_threshold = min_threshold + self.color = color + + def __lt__(self, other): + return self.min_threshold < other.min_threshold + + @classmethod + def from_viridis(cls, min_threshold): + cm = matplotlib.colormaps['viridis'] + color: Color = tuple((int(255 * cm(min_threshold)[i]) for i in range(3))) + return cls(min_threshold, color) + + 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. @@ -26,6 +43,10 @@ class VideoWriter(abc.ABC): def write_frame(self, frame: VideoFrame) -> None: pass + @abc.abstractmethod + def plot_boxes_to_frame(self, data: list[DetectedObject], frame: VideoFrame) -> None: + pass + @abc.abstractmethod def close(self) -> None: pass @@ -37,6 +58,9 @@ class VideoWriter(abc.ABC): class NoopVideoWriter(VideoWriter): + def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None: + pass + def write_frame(self, frame: VideoFrame) -> None: pass @@ -49,24 +73,39 @@ class NoopVideoWriter(VideoWriter): class MP4VVideoWriter(VideoWriter): - def __init__(self, video_writer: cv2.VideoWriter) -> None: + def __init__(self, video_writer: cv2.VideoWriter, thresholds: list[DetectedObjectBox]) -> None: self._video_writer = video_writer + self._thresholds = thresholds def write_frame(self, frame: VideoFrame) -> None: self._video_writer.write(frame) + def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None: + confidence = float(data.confidence) + for threshold in self._thresholds: + if threshold.min_threshold < confidence: + # xmin, ymin, xmax, ymax = self._extract_box(data) + cv2.rectangle(frame, *data.as_int_box(), threshold.color, 2) + return + + def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None: + for box in boxes: + self.plot_box_to_frame(box, frame) + + def close(self) -> None: self._video_writer.release() cv2.destroyAllWindows() def wait(self) -> bool: - return False #cv2.waitKey(1) == ord("q") + return False @classmethod - def from_frame(cls, input_: Path, output: Path) -> VideoWriter: + def from_frame(cls, input_: Path, output: Path, thresholds: list[float]) -> VideoWriter: + object_drawers = [DetectedObjectBox.from_viridis(t) for t in thresholds] video_cap = cv2.VideoCapture(str(input_)) try: - return cls(create_video_writer(video_cap, output)) + return cls(create_video_writer(video_cap, output), object_drawers) finally: video_cap.release() diff --git a/flake.nix b/flake.nix index b7e2372..359e704 100644 --- a/flake.nix +++ b/flake.nix @@ -27,6 +27,7 @@ virtual-env = pkgs.python312.withPackages (ps: [ ps.matplotlib ps.torchWithCuda + ps.jupyter pythonPackages.opencv4 pythonPackages.boxmot pythonPackages.ultralytics