Going somewhere with a glass of wine.
This commit is contained in:
1 parent
4171dd2bc4
commit
bcadb5b9f7
3 files changed
+8
-37
No files matched your search
@@ -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']
|
||||
|
||||
Reference in new issue
Block a user