Going somewhere without a glass of wine.

dodgeballtracker/logic.py

@@ -130,14 +130,20 @@ class DodgeballTracker:
         """
         load the pre-trained YOLOv8n model
         """
-        return YOLO("yolo11n.pt", verbose=False)
+        my_model = Path(
+            '/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)
     def update_tracker(self, tracker: DeepOCSORT, detected_objs: pd.DataFrame, frame: VideoFrame):
-        data_ = tracker.update(detected_objs.to_numpy(), frame)
+        data = tracker.update(detected_objs.to_numpy(), frame)
         columns = ['xmin', 'ymin', 'xmax', 'ymax', 'object_id', 'score', 'category', 'index']
-        return pd.DataFrame(data_, columns=columns)
+        if not data.size:
+            return pd.DataFrame({col: pd.Series(dtype='float64') for col in columns})
+        return pd.DataFrame(data, columns=columns)
 
+    def train(self):
+        model = self._get_detection_model()
 
     def main(self):
         # define some constants
@@ -149,7 +155,7 @@ 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.plot_boxes_to_frame(tracked_obj, frame)
             self.video_writer.write_frame(frame)
 
             if self.video_writer.wait():