diff --git a/data/.gitignore b/data/.gitignore index e69de29..c97ff52 100644 --- a/data/.gitignore +++ b/data/.gitignore @@ -0,0 +1,2 @@ +good-images-for-training +label-studio-exports diff --git a/dodgeballtracker/__main__.py b/dodgeballtracker/__main__.py index 5220857..2cc0a83 100644 --- a/dodgeballtracker/__main__.py +++ b/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) diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index 74256b9..4d4a746 100644 --- a/dodgeballtracker/logic.py +++ b/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(): diff --git a/dodgeballtracker/training_experiments.py b/dodgeballtracker/training_experiments.py new file mode 100644 index 0000000..916a180 --- /dev/null +++ b/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() + diff --git a/dodgeballtracker/typing.py b/dodgeballtracker/typing.py index 37ee5e7..d3262da 100644 --- a/dodgeballtracker/typing.py +++ b/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] diff --git a/dodgeballtracker/video_writer.py b/dodgeballtracker/video_writer.py index 1fb574c..4fcf544 100644 --- a/dodgeballtracker/video_writer.py +++ b/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():