Going somewhere without a glass of wine.

dodgeballtracker/__main__.py

@@ -60,6 +60,11 @@ def process_logic(input_: Path, output: Path | None, display_video: bool):
         with time_tracker():
             tracker.main()
 
+
+@cli.command()
+def train():
+    pass
+
 @cli.command()
 @click.option('--input', '-i', 'input_',
               type=click.Path(readable=True, path_type=Path))
@@ -71,4 +76,4 @@ def process(input_: Path, output: Path | None, display_video: bool):
 
 if __name__ == "__main__":
     cli()
-    # process_logic(Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/micro.mp4'), None, False)
+    # process_logic(Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/experiments/rue-honore-beaugrand-ouest-001.MOV'), None, False)

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():

dodgeballtracker/training_experiments.py

@@ -0,0 +1,35 @@
+from ultralytics import YOLO
+from pathlib import Path
+
+
+def foo():
+    model = YOLO("yolo11x.pt")
+    example = Path(
+        '/home/didier/all-dressed-programming/dodgeball-tracker/data/good-images-for-training/PB120196.JPG')
+    results = model.predict(example)
+    print('print done predicting')
+    print(f'done predicting; {len(results)} results')
+    for res in results:
+        print(f"    >>> {res.boxes.data}")
+
+def main():
+    # Load a model
+    model = YOLO("yolo11s.pt")  # load a pretrained model (recommended for training)
+
+    # Train the model
+    train = '/home/didier/all-dressed-programming/dodgeball-tracker/data/wheels_004.yaml'
+    model.train(data=Path(train),
+                epochs=150,
+                imgsz=640,
+                cache='ram')
+    print('Done training')
+    bad_example = Path("/home/didier/all-dressed-programming/dodgeball-tracker/data/label-studio-exports/2024-11-12_001/images/train/2a6267f7-PB120166.JPG")
+    example = Path('/home/didier/all-dressed-programming/dodgeball-tracker/data/good-images-for-training/PB120196.JPG')
+    results = model.predict(bad_example)
+    print(f'done predicting; {len(results)} results')
+    for res in results:
+        print(f"    >>> {res.boxes.data}")
+
+if __name__ == '__main__':
+    foo()
+

dodgeballtracker/typing.py

@@ -1,6 +1,4 @@
 from typing import NewType
-from dataclasses import dataclass
-import numpy as np
 import pandera as pa
 from pandera.dtypes import Float32, Float64
 
@@ -17,7 +15,6 @@ DetectedObject_ = tensor_of(Float32, ["xmin", "ymin", "xmax", "ymax", "score",
 
 TrackedDetectedObject = tensor_of(Float64,
                                   list(DetectedObject_.columns) + ["object_id", "index"])
-# ["xmin", "ymin", "xmax", "ymax", "object_id", "score", "category", "index"])
 
 
 Color = tuple[int, int, int]

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, Color, DetectedObject_
+from dodgeballtracker.typing import (VideoFrame, Color, TrackedDetectedObject)
 import matplotlib
 import cv2
 
@@ -84,17 +84,34 @@ class MP4VVideoWriter(VideoWriter):
     def write_frame(self, frame: VideoFrame) -> None:
         self._video_writer.write(frame)
 
+    def write_box_id(self, frame: VideoFrame, box: pd.Series):
+        bottom_left_corner_of_text = (int(box['xmin']), int(box['ymax']))
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        font_scale = 0.5
+        font_color = (255, 255, 255)
+        thickness = 1
+        line_type = 2
+        id_ = str(int(box['object_id']))
+        cv2.putText(frame,
+                    id_,
+                    bottom_left_corner_of_text,
+                    font,
+                    font_scale,
+                    font_color,
+                    thickness,
+                    line_type)
+
     def plot_box_to_frame(self, data: pd.Series, frame: VideoFrame) -> None:
-        confidence = float(data.score)
-        # data = data.map(int)
+        confidence = data['score']
         for threshold in self._thresholds:
             if threshold.min_threshold < confidence:
                 start_point =  data[['xmin', 'ymin']].array.astype('int32')
                 end_point = data[['xmax', 'ymax']].array.astype('int32')
                 cv2.rectangle(frame, start_point, end_point, threshold.color, 2)
+                self.write_box_id(frame, data)
                 return
 
-    @pa.decorators.check_io(boxes=DetectedObject_)
+    @pa.decorators.check_io(boxes=TrackedDetectedObject)
     def plot_boxes_to_frame(self, boxes: pd.DataFrame, frame: VideoFrame) -> None:
         boxes_that_passe_threshold = boxes[boxes['score'] > 0.5]
         for i, box in boxes_that_passe_threshold.iterrows():