From 3ea54208bec15638c5e6e1e4958f2275089cdbd1 Mon Sep 17 00:00:00 2001 From: didier amyot Date: Tue, 23 Jan 2024 16:12:53 -0500 Subject: [PATCH] migrated to class --- dodgeballtracker/__init__.py | 7 +- dodgeballtracker/logic.py | 144 ++++++++++++++++++----------------- 2 files changed, 77 insertions(+), 74 deletions(-) diff --git a/dodgeballtracker/__init__.py b/dodgeballtracker/__init__.py index dc03919..6fbfc2a 100644 --- a/dodgeballtracker/__init__.py +++ b/dodgeballtracker/__init__.py @@ -1,6 +1,6 @@ from pathlib import Path -from dodgeballtracker.logic import DetectedObjectBox, main +from dodgeballtracker.logic import DetectedObjectBox, DodgeballTracker import numpy as np @@ -12,8 +12,9 @@ if __name__ == "__main__": detected_box_objects = [ DetectedObjectBox.from_viridis(threshold) for threshold in np.arange(0.5, 1, 0.05) ] - input_video = input_dir / "micro.mp4" output_video = output_dir / "output.mp4" - main(input_video, output_video, detected_box_objects) + dodgeball_tracker = DodgeballTracker(input_video, output_video, detected_box_objects) + + dodgeball_tracker.main() diff --git a/dodgeballtracker/logic.py b/dodgeballtracker/logic.py index aab86dd..f4d259f 100644 --- a/dodgeballtracker/logic.py +++ b/dodgeballtracker/logic.py @@ -26,95 +26,97 @@ class DetectedObjectBox: return cls(min_threshold, color) -def write_frame(writer, frame): - cv2.imshow("Frame", frame) - writer.write(frame) +class DodgeballTracker: + def __init__(self, + input_file: Path, + output_file: Path, + detected_box_objects: list[DetectedObjectBox]): + self.input_path = input_file + self.output_path = output_file + self.detected_box_objects = sorted(detected_box_objects, reverse=True) -def extract_box(data): - return int(data[0]), int(data[1]), int(data[2]), int(data[3]) + def write_frame(self, writer, frame) -> None: + cv2.imshow("Frame", frame) + writer.write(frame) + def extract_box(self, data) -> tuple[int, int, int, int]: + return int(data[0]), int(data[1]), int(data[2]), int(data[3]) -def plot_boxe_to_frame(data, frame, detected_box_objects): - confidence = float(data[4]) - for detected_box_object in detected_box_objects: + def plot_boxe_to_frame(self, data, frame): + confidence = float(data[4]) + for detected_box_object in self.detected_box_objects: - if detected_box_object.min_threshold < confidence: - xmin, ymin, xmax, ymax = extract_box(data) - cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2) - return + if detected_box_object.min_threshold < confidence: + xmin, ymin, xmax, ymax = self.extract_box(data) + cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), detected_box_object.color, 2) + return + def run_model_on_frame(self, model, frame): + return model(frame)[0] -def run_model_on_frame(model, frame): - return model(frame)[0] + def get_court_perspective_transform(self): + """ + """ + empirical_top_left = np.float32([416, 175]) + empirical_top_right = np.float32([905, 173]) + empirical_bottom_left = np.float32([51, 702]) + empirical_bottom_right = np.float32([1188, 695]) + court_top_left = np.float32([0, 0]) + court_top_right = np.float32([9, 0]) + court_bottom_left = np.float32([0, 18]) + court_bottom_right = np.float32([9, 18]) -def crazy_experiment(frame): - cv2.putText(frame, - "Hello world", - (50, 50), - cv2.FONT_HERSHEY_SIMPLEX, - 2, - (0, 255, 0), 8) - cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5) - pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32) - pts = pts.reshape((-1, 1, 2)) + empirical_source = np.float32([empirical_top_left, + empirical_top_right, + empirical_bottom_left, + empirical_bottom_right]) + pixel_per_meter = 50 + empirical_target = np.float32([court_top_left, + court_top_right, + court_bottom_left, + court_bottom_right]) * pixel_per_meter + return cv2.getPerspectiveTransform(empirical_source, empirical_target) - cv2.polylines(frame, pts, True, (0, 255, 255), 8) + def crazy_experiment(self, frame): + cv2.putText(frame, + "Hello world", + (50, 50), + cv2.FONT_HERSHEY_SIMPLEX, + 2, + (0, 255, 0), 8) + cv2.line(frame, (0, 0), (511, 511), (255, 0, 0), 5) + pts = np.array([[50, 50], [500, 50], [500, 500], [50, 500], [50, 50]], np.int32) + pts = pts.reshape((-1, 1, 2)) + cv2.polylines(frame, pts, True, (0, 255, 255), 8) -def get_court_perspective_transform(): - """ - """ - empirical_top_left = np.float32([416, 175]) - empirical_top_right = np.float32([905, 173]) - empirical_bottom_left = np.float32([51, 702]) - empirical_bottom_right = np.float32([1188, 695]) + def main(self): + # define some constants - court_top_left = np.float32([0, 0]) - court_top_right = np.float32([9, 0]) - court_bottom_left = np.float32([0, 18]) - court_bottom_right = np.float32([9, 18]) + video_cap = cv2.VideoCapture(str(self.input_path)) + writer = create_video_writer(video_cap, str(self.output_path)) - empirical_source = np.float32([empirical_top_left, - empirical_top_right, - empirical_bottom_left, - empirical_bottom_right]) - pixel_per_meter = 50 - empirical_target = np.float32([court_top_left, - court_top_right, - court_bottom_left, - court_bottom_right]) * pixel_per_meter - return cv2.getPerspectiveTransform(empirical_source, empirical_target) + # load the pre-trained YOLOv8n model + model = YOLO("yolov8n.pt") + for frame in iter_frames(video_cap): + with time_tracker(): + # run the YOLO model on the frame + detections = self.run_model_on_frame(model, frame) + for data in detections.boxes.data.tolist(): + self.plot_boxe_to_frame(data, frame) -def main(input_filename: Path, output_filename: Path, detected_box_objects): - # define some constants - detected_box_objects = sorted(detected_box_objects, reverse=True) + # show the frame to our screen + self.write_frame(writer, frame) - video_cap = cv2.VideoCapture(str(input_filename)) - writer = create_video_writer(video_cap, str(output_filename)) + if cv2.waitKey(1) == ord("q"): + break - # load the pre-trained YOLOv8n model - model = YOLO("yolov8n.pt") - tracker = DeepSort(max_age=50) - - for frame in iter_frames(video_cap): - with time_tracker(): - # run the YOLO model on the frame - detections = run_model_on_frame(model, frame) - for data in detections.boxes.data.tolist(): - plot_boxe_to_frame(data, frame, detected_box_objects) - - # show the frame to our screen - write_frame(writer, frame) - - if cv2.waitKey(1) == ord("q"): - break - - video_cap.release() - writer.release() - cv2.destroyAllWindows() + video_cap.release() + writer.release() + cv2.destroyAllWindows()