Going somewhere with a coffee.

This commit is contained in:
didier committed 2024-11-15 10:00:22 -05:00
1 parent 84b9c2674c
commit f5b8e24e01
6 files changed
+191 -5

No files matched your search

+3 -3
View File
@@ -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__)
@@ -107,7 +107,7 @@ class DodgeballTracker:
return True
return False
@pa.check_output(DetectedObject_)
@pa.check_output(DetectedObject)
def run_model_on_frame(self, model, frame: VideoFrame):
model_outputs = model(frame)[0]
boxes = model_outputs.boxes.data.cpu().numpy()
@@ -134,7 +134,7 @@ class DodgeballTracker:
'/home/didier/all-dressed-programming/dodgeball-tracker/dodgeballtracker/runs/detect/train31/weights/best.pt')
return YOLO(my_model, 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']