| from dataclasses import dataclass |
|
|
| import cv2 |
|
|
| from src.methods.detection.base import Detection |
| from src.methods.detection.mot import MOT17Detector |
| from src.methods.tracking import SORT |
| from src.methods.tracking.base import Track |
| from src.methods.tracking.mot_ground_truth import MOTGroundTruthTracker |
| from src.utils.mot17 import resolve_mot17_sequence_dir |
| from src.utils.render import render_detections, render_tracks |
|
|
| |
| DETECTORS = { |
| "mot17": MOT17Detector, |
| } |
|
|
| TRACKERS = { |
| "sort": SORT, |
| "ground_truth": MOTGroundTruthTracker, |
| } |
|
|
| @dataclass |
| class FrameOutputs: |
| detections: list[Detection] |
| tracks: list[Track] |
|
|
| class FrameDebugger: |
| """Utility for inspecting and debugging frame-wise inputs and outputs of MOT methods.""" |
|
|
| @property |
| def frame_count(self) -> int: |
| return len(self.frame_paths) |
|
|
| def __init__(self, mot_seq: str, detector_name: str | None, tracker_name: str | None): |
| self.frame_outputs = [] |
| self.sequence_dir = resolve_mot17_sequence_dir(mot_seq) |
| self.mot_seq = self.sequence_dir.name |
| self.frame_paths = sorted((self.sequence_dir / "img1").glob("*.jpg")) |
| self.frame_outputs: list[FrameOutputs] = [] |
|
|
| self.detector = None |
| if detector_name is not None: |
| detector_cls = DETECTORS[detector_name] |
| self.detector = detector_cls( |
| sequence_id=self.mot_seq, |
| root_dir=self.sequence_dir.parent, |
| ) |
| |
| self.tracker = None |
| if tracker_name is not None: |
| tracker_cls = TRACKERS[tracker_name] |
| self.tracker = tracker_cls( |
| sequence_id=self.mot_seq, |
| root_dir=self.sequence_dir.parent, |
| ) if tracker_name == "ground_truth" else tracker_cls() |
|
|
| |
| self._run() |
| |
| def _run(self) -> None: |
| """ |
| Run tracking and detection for all the frames in the sequence and store the outputs for later. |
| Safe to run again but should not be needed ever. |
| """ |
| self.frame_outputs = [] |
| for frame_index, _ in enumerate(self.frame_paths, start=1): |
| detections = [] |
| if self.detector is not None: |
| detections = self.detector.get_detections(frame_index=frame_index) |
|
|
| tracks = [] |
| if self.tracker is not None: |
| tracks = self.tracker.update(detections, frame_index=frame_index) |
|
|
| self.frame_outputs.append( |
| FrameOutputs( |
| detections=detections, |
| tracks=tracks, |
| ) |
| ) |
|
|
| def get_annotated_frame(self, frame_index: int) -> tuple[object | None, str]: |
| """ |
| Get the frame at the index with annotations rendered for the detector or tracker if they were provided. |
| Returns a tuple of (annotated_frame, status_message). |
| If the frame index is out of bounds or the frame cannot be read, returns (None, error_message). |
| """ |
| if frame_index < 1 or frame_index > self.frame_count: |
| return None, f"Frame index {frame_index} is out of bounds. Expected 1 to {self.frame_count}." |
|
|
| frame_path = self.frame_paths[frame_index - 1] |
| frame = cv2.imread(str(frame_path)) |
| if frame is None: |
| return None, f"Could not read frame {frame_index}: {frame_path}" |
|
|
| outputs = self.frame_outputs[frame_index - 1] |
| annotated_frame = frame |
|
|
| if outputs.tracks: |
| annotated_frame = render_tracks(annotated_frame, outputs.tracks) |
| elif outputs.detections: |
| annotated_frame = render_detections(annotated_frame, outputs.detections) |
|
|
| |
|
|
| annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB) |
|
|
| status = ( |
| f"Frame {frame_index} / {self.frame_count} " |
| f"| detections: {len(outputs.detections)} " |
| f"| tracks: {len(outputs.tracks)}" |
| ) |
| return annotated_frame, status |
|
|
|
|
|
|
