diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index bf02118..74256b9 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_, TrackedDetectedObject_ +from .typing import VideoFrame, DetectedObject_, TrackedDetectedObject logger = logging.getLogger(__name__) @@ -112,7 +112,8 @@ class DodgeballTracker: model_outputs = model(frame)[0] boxes = model_outputs.boxes.data.cpu().numpy() columns=['xmin', 'ymin', 'xmax', 'ymax', 'score', 'category'] - return pd.DataFrame(boxes, columns=columns) + ret = pd.DataFrame(boxes, columns=columns) + return ret def persist_points(self): with open(self.output_json_path, "wt") as file_: @@ -129,9 +130,9 @@ class DodgeballTracker: """ load the pre-trained YOLOv8n model """ - return YOLO("yolo11n.pt", verbose=True) + return YOLO("yolo11n.pt", verbose=False) - @pa.check_io(detected_objs=DetectedObject_, out=TrackedDetectedObject_) + @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', 'ymax', 'object_id', 'score', 'category', 'index'] diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index 9cfb6db..37ee5e7 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -4,40 +4,10 @@ import numpy as np import pandera as pa from pandera.dtypes import Float32, Float64 -XCoord = NewType('XCoord', float) -YCoord = NewType('YCoord', float) -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 - ymin: YCoord - xmax: XCoord - ymax: YCoord - score: Score - 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.score, - self.category]) - - def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema: return pa.DataFrameSchema({col: pa.Column(data_type) for col in cols}) @@ -45,8 +15,8 @@ def tensor_of(data_type: pa.DataType, cols: list[str]) -> pa.DataFrameSchema: DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score", "category"]) -TrackedDetectedObject_ = tensor_of(Float64, - list(DetectedObject_.columns) + ["object_id", "index"]) +TrackedDetectedObject = tensor_of(Float64, + list(DetectedObject_.columns) + ["object_id", "index"]) # ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"]) diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py index b6866d1..1fb574c 100644 --- a/dodgeballtracker/video_writer.py +++ b/dodgeballtracker/video_writer.py @@ -4,7 +4,7 @@ from pathlib import Path import pandas as pd import pandera as pa -from dodgeballtracker.typing import VideoFrame, DetectedObject, Color, DetectedObject_ +from dodgeballtracker.typing import VideoFrame, Color, DetectedObject_ import matplotlib import cv2