diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index b40b1e5..c9c6f0d 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -7,8 +7,10 @@ import cv2 from dodgeballtracker.helper import iter_frames from dodgeballtracker.video_writer import VideoWriter, DetectedObjectBox 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_ logger = logging.getLogger(__name__) @@ -105,21 +107,18 @@ class DodgeballTracker: return True return False + @pa.check_output(DetectedObject_) def run_model_on_frame(self, model, frame: VideoFrame): model_outputs = model(frame)[0] - return model_outputs.boxes.data.cpu().numpy() - # return model_outputs.boxes.data.tolist() - foo = model_outputs.boxes.data - for data in model_outputs.boxes.data.tolist(): - detected_obj = DetectedObject(*data) - if detected_obj.score > 0.5: - yield detected_obj + boxes = model_outputs.boxes.data.cpu().numpy() + columns=['xmin', 'ymin', 'xmax', 'ymin', 'score', 'category'] + return pd.DataFrame(boxes, columns=columns) def persist_points(self): with open(self.output_json_path, "wt") as file_: json.dump({"points": self.points}, file_, indent=4) - def _get_tracker(self): + def _get_tracker(self) -> DeepOCSORT: return DeepOCSORT( Path('lmbn_n_cuhk03_d.pt'), # which ReID model to use get_best_available_torch_device(), @@ -132,8 +131,11 @@ 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.score > 0.5] + def update_tracker(self, tracker: DeepOCSORT, detected_objs, frame: VideoFrame): + data_ = tracker.update(detected_objs.to_numpy(), frame) + columns = ['xmin', 'ymin', 'xmax', 'ymin', 'object_id', 'score', 'category', 'index'] + return pd.DataFrame(data_, columns=columns) + def main(self): # define some constants @@ -143,18 +145,11 @@ class DodgeballTracker: for i, (frame, ts) in enumerate(iter_frames(video_cap)): # run the YOLO model on the frame - detected_objs = list(self.run_model_on_frame(model, frame)) - data_ = tracker.update( - np.array(detected_objs), frame) - if i == 100: - print(data_) + 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) - # 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) + # + # self.video_writer.write_frame(frame) if self.video_writer.wait(): break diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index d8612e5..5a9a71a 100644 --- a/dodgeballtracker/typing.py +++ b/dodgeballtracker/typing.py @@ -1,6 +1,8 @@ from typing import NewType from dataclasses import dataclass import numpy as np +import pandera as pa +from pandera.dtypes import Float32 XCoord = NewType('XCoord', float) YCoord = NewType('YCoord', float) @@ -33,5 +35,25 @@ class DetectedObject: 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), +}) + Color = tuple[int, int, int] \ No newline at end of file