diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index 10762c2..5220857 100644 --- a/dodgeballtracker/__main__.py +++ b/dodgeballtracker/__main__.py @@ -41,12 +41,7 @@ def get_video_writer(display_frame: bool, 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): +def process_logic(input_: Path, output: Path | None, display_video: bool): output = get_output(input_, output) detected_box_objects = [ DetectedObjectBox.from_viridis(threshold) for threshold in @@ -65,6 +60,15 @@ def process(input_: Path, output: Path | None, display_video: bool): with time_tracker(): tracker.main() +@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): + process_logic(input_, output, display_video) + if __name__ == "__main__": cli() + # process_logic(Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/micro.mp4'), None, False) diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index c9c6f0d..bf02118 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -10,7 +10,7 @@ import numpy as np import pandas as pd import pandera as pa from boxmot import DeepOcSort as DeepOCSORT -from .typing import VideoFrame, DetectedObject_ +from .typing import VideoFrame, DetectedObject_, TrackedDetectedObject_ logger = logging.getLogger(__name__) @@ -111,7 +111,7 @@ class DodgeballTracker: def run_model_on_frame(self, model, frame: VideoFrame): model_outputs = model(frame)[0] boxes = model_outputs.boxes.data.cpu().numpy() - columns=['xmin', 'ymin', 'xmax', 'ymin', 'score', 'category'] + columns=['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category'] return pd.DataFrame(boxes, columns=columns) def persist_points(self): @@ -131,9 +131,10 @@ class DodgeballTracker: """ return YOLO("yolo11n.pt", verbose=True) - def update_tracker(self, tracker: DeepOCSORT, detected_objs, frame: VideoFrame): + @pa.check_io(detected_objs=DetectedObject_, out=TrackedDetectedObject_) + def update_tracker(self, tracker: DeepOCSORT, detected_objs: pd.DataFrame, frame: VideoFrame): data_ = tracker.update(detected_objs.to_numpy(), frame) - columns = ['xmin', 'ymin', 'xmax', 'ymin', 'object_id', 'score', 'category', 'index'] + columns = ['xmin', 'ymin', 'xmax', 'ymax', 'object_id', 'score', 'category', 'index'] return pd.DataFrame(data_, columns=columns) @@ -147,9 +148,8 @@ class DodgeballTracker: # run the YOLO model on the frame detected_objs = self.run_model_on_frame(model, frame) tracked_obj = self.update_tracker(tracker, detected_objs, frame) - # self.video_writer.plot_boxes_to_frame(detected_objs, frame) - # - # self.video_writer.write_frame(frame) + self.video_writer.plot_boxes_to_frame(detected_objs, frame) + self.video_writer.write_frame(frame) if self.video_writer.wait(): break diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index 5a9a71a..9cfb6db 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -2,7 +2,7 @@ from typing import NewType from dataclasses import dataclass import numpy as np import pandera as pa -from pandera.dtypes import Float32 +from pandera.dtypes import Float32, Float64 XCoord = NewType('XCoord', float) YCoord = NewType('YCoord', float) @@ -10,15 +10,16 @@ Score = 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 @@ -29,31 +30,24 @@ class DetectedObject: 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)) + 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.score, self.category]) + return np.array([self.xmin, self.ymax, self.xmax, self.ymax, self.score, + self.category]) -DetectedObject_ = pa.DataFrameSchema({ - "xmin": pa.Column(Float32), - "ymin": pa.Column(Float32), - "xmax": pa.Column(Float32), - "ymin": pa.Column(Float32), - "score": pa.Column(Float32), - "category": pa.Column(Float32), -}) - -TrackedDetectedObject_ = pa.DataFrameSchema({ - "xmin": pa.Column(Float32), - "ymin": pa.Column(Float32), - "xmax": pa.Column(Float32), - "ymin": pa.Column(Float32), - "object_id": pa.Column(Float32), - "score": pa.Column(Float32), - "category": pa.Column(Float32), - "index": pa.Column(Float32), -}) +def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema: + return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols}) -Color = tuple[int, int, int] \ No newline at end of file +DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score", + "category"]) + +TrackedDetectedObject_ = tensor_of(Float64, + list(DetectedObject_.columns) + ["object_id", "index"]) +# ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"]) + + +Color = tuple[int, int, int] diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py index 0c0273f..b6866d1 100644 --- a/dodgeballtracker/video_writer.py +++ b/dodgeballtracker/video_writer.py @@ -1,6 +1,10 @@ import abc from pathlib import Path -from dodgeballtracker.typing import VideoFrame, DetectedObject, Color + +import pandas as pd +import pandera as pa + +from dodgeballtracker.typing import VideoFrame, DetectedObject, Color, DetectedObject_ import matplotlib import cv2 @@ -44,7 +48,7 @@ class VideoWriter(abc.ABC): pass @abc.abstractmethod - def plot_boxes_to_frame(self, data: list[DetectedObject], frame: VideoFrame) -> None: + def plot_boxes_to_frame(self, data: pd.DataFrame, frame: VideoFrame) -> None: pass @abc.abstractmethod @@ -58,7 +62,7 @@ class VideoWriter(abc.ABC): class NoopVideoWriter(VideoWriter): - def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None: + def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None: pass def write_frame(self, frame: VideoFrame) -> None: @@ -80,16 +84,20 @@ class MP4VVideoWriter(VideoWriter): def write_frame(self, frame: VideoFrame) -> None: self._video_writer.write(frame) - def plot_box_to_frame(self, data: DetectedObject, frame: VideoFrame) -> None: + def plot_box_to_frame(self, data: pd.Series, frame: VideoFrame) -> None: confidence = float(data.score) + # data = data.map(int) 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) + start_point = data[['xmin', 'ymin']].array.astype('int32') + end_point = data[['xmax', 'ymax']].array.astype('int32') + cv2.rectangle(frame, start_point, end_point, threshold.color, 2) return - def plot_boxes_to_frame(self, boxes: list[DetectedObject], frame: VideoFrame) -> None: - for box in boxes: + @pa.decorators.check_io(boxes=DetectedObject_) + def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None: + boxes_that_passe_threshold = boxes[boxes['score'] > 0.5] + for i, box in boxes_that_passe_threshold.iterrows(): self.plot_box_to_frame(box, frame)