supervisely 6.73.410__py3-none-any.whl → 6.73.470__py3-none-any.whl

supervisely/nn/tracker/botsort_tracker.py

@@ -0,0 +1,273 @@
+import supervisely as sly
+from supervisely.nn.tracker.base_tracker import BaseTracker
+from supervisely import Annotation, VideoAnnotation
+from supervisely.annotation.label import LabelingStatus
+from dataclasses import dataclass
+from types import SimpleNamespace
+from typing import List, Dict, Tuple, Any, Optional
+import numpy as np
+import yaml
+import os
+from pathlib import Path
+from supervisely import logger
+from supervisely.nn.tracker.botsort.tracker.mc_bot_sort import BoTSORT
+
+
+@dataclass
+class TrackedObject:
+    """
+    Data class representing a tracked object in a single frame.
+    
+    Args:
+        track_id: Unique identifier for the track
+        det_id: Detection ID for mapping back to original annotation
+        bbox: Bounding box coordinates in format [x1, y1, x2, y2]
+        class_name: String class name
+        class_sly_id: Supervisely class ID (from ObjClass.sly_id)
+        score: Confidence score of the detection/track
+    """
+    track_id: int
+    det_id: int
+    bbox: List[float]  # [x1, y1, x2, y2]
+    class_name: str
+    class_sly_id: Optional[int]  # Supervisely class ID
+    score: float
+
+
+class BotSortTracker(BaseTracker):
+    
+    def __init__(self, settings: dict = None, device: str = None):
+        super().__init__(settings=settings, device=device)
+
+        from supervisely.nn.tracker import TRACKING_LIBS_INSTALLED
+        if not TRACKING_LIBS_INSTALLED:
+            raise ImportError(
+                "Tracking dependencies are not installed. "
+                "Please install supervisely with `pip install supervisely[tracking]`."
+            )
+        
+        # Load default settings from YAML file
+        self.settings = self._load_default_settings()
+        
+        # Override with user settings if provided
+        if settings:
+            self.settings.update(settings)
+        
+        args = SimpleNamespace(**self.settings)
+        args.name = "BotSORT"
+        args.device = self.device
+            
+        self.tracker = BoTSORT(args=args)
+        
+        # State for accumulating results
+        self.frame_tracks = []
+        self.obj_classes = {}   # class_id -> ObjClass
+        self.current_frame = 0
+        self.class_ids = {}  # class_name -> class_id mapping
+        self.frame_shape = () 
+
+    def _load_default_settings(self) -> dict:
+        """Internal method: calls classmethod"""
+        return self.get_default_params()
+
+    def update(self, frame: np.ndarray, annotation: Annotation) -> List[Dict[str, Any]]:
+        """Update tracker and return list of matches for current frame."""
+        self.frame_shape = frame.shape[:2]
+        self._update_obj_classes(annotation)
+        detections = self._convert_annotation(annotation)
+        output_stracks, detection_track_map = self.tracker.update(detections, frame)
+        tracks = self._stracks_to_tracks(output_stracks, detection_track_map)
+        
+        # Store tracks for VideoAnnotation creation
+        self.frame_tracks.append(tracks)
+        self.current_frame += 1
+        
+        matches = []
+        for pair in detection_track_map:
+            det_id = pair["det_id"]
+            track_id = pair["track_id"]
+
+            if track_id is not None:
+                match = {
+                    "track_id": track_id,
+                    "label": annotation.labels[det_id]
+                }
+                matches.append(match)
+            
+        return matches
+    
+    def reset(self) -> None:
+        super().reset()
+        self.frame_tracks = []
+        self.obj_classes = {}
+        self.current_frame = 0
+        self.class_ids = {}
+        self.frame_shape = ()
+
+    def track(self, frames: List[np.ndarray], annotations: List[Annotation]) -> VideoAnnotation:
+        """Track objects through sequence of frames and return VideoAnnotation."""
+        if len(frames) != len(annotations):
+            raise ValueError("Number of frames and annotations must match")
+        
+        self.reset()
+        
+        # Process each frame
+        for frame_idx, (frame, annotation) in enumerate(zip(frames, annotations)):
+            self.current_frame = frame_idx
+            self.update(frame, annotation)
+        
+        # Convert accumulated tracks to VideoAnnotation
+        return self._create_video_annotation()
+        
+    def _convert_annotation(self, annotation: Annotation) -> np.ndarray:
+        """Convert Supervisely annotation to BoTSORT detection format."""
+        detections_list = []
+
+        for label in annotation.labels:
+            if label.tags.get("confidence", None) is not None:
+                confidence = label.tags.get("confidence").value
+            elif label.tags.get("conf", None) is not None:
+                confidence = label.tags.get("conf").value
+            else:
+                confidence = 1.0
+                logger.debug(
+                    f"Label {label.obj_class.name} does not have confidence tag, using default value 1.0"
+                )
+
+            rectangle = label.geometry.to_bbox()
+            
+            class_name = label.obj_class.name
+            class_id = self.class_ids[class_name]
+            
+            detection = [
+                rectangle.left,    # x1
+                rectangle.top,     # y1 
+                rectangle.right,   # x2
+                rectangle.bottom,  # y2   
+                confidence,        # score
+                class_id,          # class_id as number
+            ]
+            detections_list.append(detection)
+
+        if detections_list:
+            return np.array(detections_list, dtype=np.float32)
+        else:
+            return np.zeros((0, 6), dtype=np.float32)
+    
+    def _stracks_to_tracks(self, output_stracks, detection_track_map) -> List[TrackedObject]:
+        """Convert BoTSORT output tracks to TrackedObject dataclass instances."""
+        tracks = []
+        
+        id_to_name = {v: k for k, v in self.class_ids.items()}
+        
+        track_id_to_det_id = {}
+        for pair in detection_track_map:
+            det_id = pair["det_id"]
+            track_id = pair["track_id"]
+            track_id_to_det_id[track_id] = det_id 
+        
+        for strack in output_stracks:
+            # BoTSORT may store class info in different attributes
+            # Try to get class_id from various possible sources
+            class_id = 0  # default
+            
+            if hasattr(strack, 'cls') and strack.cls != -1:
+                # cls should contain the numeric ID we passed in
+                class_id = int(strack.cls)
+            elif hasattr(strack, 'class_id'):
+                class_id = int(strack.class_id)
+            
+            class_name = id_to_name.get(class_id, "unknown")
+            
+            # Get Supervisely class ID from stored ObjClass
+            class_sly_id = None
+            if class_name in self.obj_classes:
+                obj_class = self.obj_classes[class_name]
+                class_sly_id = obj_class.sly_id
+            
+            track = TrackedObject(
+                track_id=strack.track_id,
+                det_id=track_id_to_det_id.get(strack.track_id),
+                bbox=strack.tlbr.tolist(),  # [x1, y1, x2, y2]
+                class_name=class_name,
+                class_sly_id=class_sly_id,
+                score=getattr(strack, 'score', 1.0)
+            )
+            tracks.append(track)
+        
+        return tracks
+        
+    def _update_obj_classes(self, annotation: Annotation):
+        """Extract and store object classes from annotation."""
+        for label in annotation.labels:
+            class_name = label.obj_class.name
+            if class_name not in self.obj_classes:
+                self.obj_classes[class_name] = label.obj_class
+        
+            if class_name not in self.class_ids:
+                self.class_ids[class_name] = len(self.class_ids)
+
+
+    def _create_video_annotation(self) -> VideoAnnotation:
+        """Convert accumulated tracking results to Supervisely VideoAnnotation."""
+        img_h, img_w = self.frame_shape
+        video_objects = {}  # track_id -> VideoObject
+        frames = []
+        
+        for frame_idx, tracks in enumerate(self.frame_tracks):
+            frame_figures = []
+            
+            for track in tracks:
+                track_id = track.track_id
+                bbox = track.bbox  # [x1, y1, x2, y2]
+                class_name = track.class_name
+                
+                # Clip bbox to image boundaries
+                x1, y1, x2, y2 = bbox
+                dims = np.array([img_w, img_h, img_w, img_h]) - 1
+                x1, y1, x2, y2 = np.clip([x1, y1, x2, y2], 0, dims)
+                
+                # Get or create VideoObject
+                if track_id not in video_objects:
+                    obj_class = self.obj_classes.get(class_name)
+                    if obj_class is None:
+                        continue  # Skip if class not found
+                    video_objects[track_id] = sly.VideoObject(obj_class)
+                
+                video_object = video_objects[track_id]
+                rect = sly.Rectangle(top=y1, left=x1, bottom=y2, right=x2)
+                frame_figures.append(sly.VideoFigure(video_object, rect, frame_idx, track_id=str(track_id), status=LabelingStatus.AUTO))
+            
+            frames.append(sly.Frame(frame_idx, frame_figures))
+
+        objects = list(video_objects.values())
+
+        
+        return VideoAnnotation(
+            img_size=self.frame_shape,
+            frames_count=len(self.frame_tracks),
+            objects=sly.VideoObjectCollection(objects),
+            frames=sly.FrameCollection(frames)
+        )
+        
+    @property    
+    def video_annotation(self) -> VideoAnnotation:
+        """Return the accumulated VideoAnnotation."""
+        if not self.frame_tracks:
+            error_msg = (
+                "No tracking data available. "
+                "Please run tracking first using track() method or process frames with update()."
+            )
+            raise ValueError(error_msg)
+                
+        return self._create_video_annotation()
+
+    @classmethod
+    def get_default_params(cls) -> Dict[str, Any]:
+        """Public API: get default params WITHOUT creating instance."""
+        current_dir = Path(__file__).parent
+        config_path = current_dir / "botsort/botsort_config.yaml"
+        
+        with open(config_path, 'r', encoding='utf-8') as file:
+            return yaml.safe_load(file)
+

supervisely/nn/tracker/calculate_metrics.py

@@ -0,0 +1,264 @@
+import numpy as np
+from collections import defaultdict
+from typing import Dict, List, Union
+
+from scipy.optimize import linear_sum_assignment  # pylint: disable=import-error
+
+import supervisely as sly
+from supervisely.video_annotation.video_annotation import VideoAnnotation
+
+import motmetrics as mm  # pylint: disable=import-error
+
+class TrackingEvaluator:
+    """
+    Evaluator for video tracking metrics including MOTA, MOTP, IDF1.
+    """
+
+    def __init__(self, iou_threshold: float = 0.5):
+        """Initialize evaluator with IoU threshold for matching."""
+        from supervisely.nn.tracker import TRACKING_LIBS_INSTALLED
+        if not TRACKING_LIBS_INSTALLED:
+            raise ImportError(
+                "Tracking dependencies are not installed. "
+                "Please install supervisely with `pip install supervisely[tracking]`."
+            )
+            
+        if not 0.0 <= iou_threshold <= 1.0:
+            raise ValueError("iou_threshold must be in [0.0, 1.0]")
+        self.iou_threshold = iou_threshold
+
+    def evaluate(
+        self,
+        gt_annotation: VideoAnnotation,
+        pred_annotation: VideoAnnotation,
+    ) -> Dict[str, Union[float, int]]:
+        """Main entry: extract tracks from annotations, compute basic and MOT metrics, return results."""
+        self._validate_annotations(gt_annotation, pred_annotation)
+        self.img_height, self.img_width = gt_annotation.img_size
+
+        gt_tracks = self._extract_tracks(gt_annotation)
+        pred_tracks = self._extract_tracks(pred_annotation)
+
+        basic = self._compute_basic_metrics(gt_tracks, pred_tracks)
+        mot = self._compute_mot_metrics(gt_tracks, pred_tracks)
+
+        results = {
+            # basic detection
+            "precision": basic["precision"],
+            "recall": basic["recall"],
+            "f1": basic["f1"],
+            "avg_iou": basic["avg_iou"],
+            "true_positives": basic["tp"],
+            "false_positives": basic["fp"],
+            "false_negatives": basic["fn"],
+            "total_gt_objects": basic["total_gt"],
+            "total_pred_objects": basic["total_pred"],
+
+            # motmetrics
+            "mota": mot["mota"],
+            "motp": mot["motp"],
+            "idf1": mot["idf1"],
+            "id_switches": mot["id_switches"],
+            "fragmentations": mot["fragmentations"],
+            "num_misses": mot["num_misses"],
+            "num_false_positives": mot["num_false_positives"],
+
+            # config
+            "iou_threshold": self.iou_threshold,
+        }
+        return results
+
+    def _validate_annotations(self, gt: VideoAnnotation, pred: VideoAnnotation):
+        """Minimal type validation for annotations."""
+        if not isinstance(gt, VideoAnnotation) or not isinstance(pred, VideoAnnotation):
+            raise TypeError("gt_annotation and pred_annotation must be VideoAnnotation instances")
+
+    def _extract_tracks(self, annotation: VideoAnnotation) -> Dict[int, List[Dict]]:
+        """
+        Extract tracks from a VideoAnnotation into a dict keyed by frame index.
+        Each element is a dict: {'track_id': int, 'bbox': [x1,y1,x2,y2], 'confidence': float, 'class_name': str}
+        """
+        frames_to_tracks = defaultdict(list)
+
+        for frame in annotation.frames:
+            frame_idx = frame.index
+            for figure in frame.figures:
+                # use track_id if present, otherwise fallback to object's key int
+                track_id = int(figure.track_id) if figure.track_id is not None else figure.video_object.key().int
+
+                bbox = figure.geometry
+                if not isinstance(bbox, sly.Rectangle):
+                    bbox = bbox.to_bbox()
+
+                x1 = float(bbox.left)
+                y1 = float(bbox.top)
+                x2 = float(bbox.right)
+                y2 = float(bbox.bottom)
+
+                frames_to_tracks[frame_idx].append({
+                    "track_id": track_id,
+                    "bbox": [x1, y1, x2, y2],
+                    "confidence": float(getattr(figure, "confidence", 1.0)),
+                    "class_name": figure.video_object.obj_class.name
+                })
+
+        return dict(frames_to_tracks)
+
+    def _compute_basic_metrics(self, gt_tracks: Dict[int, List[Dict]], pred_tracks: Dict[int, List[Dict]]):
+        """
+        Compute per-frame true positives / false positives / false negatives and average IoU.
+        Matching is performed with Hungarian algorithm (scipy). Matches with IoU < threshold are discarded.
+        """
+        tp = fp = fn = 0
+        total_iou = 0.0
+        iou_count = 0
+
+        frames = sorted(set(list(gt_tracks.keys()) + list(pred_tracks.keys())))
+        for f in frames:
+            gts = gt_tracks.get(f, [])
+            preds = pred_tracks.get(f, [])
+
+            if not gts and not preds:
+                continue
+            if not gts:
+                fp += len(preds)
+                continue
+            if not preds:
+                fn += len(gts)
+                continue
+
+            gt_boxes = np.array([g["bbox"] for g in gts])
+            pred_boxes = np.array([p["bbox"] for p in preds])
+
+            # get cost matrix from motmetrics (cost = 1 - IoU)
+            cost_mat = mm.distances.iou_matrix(gt_boxes, pred_boxes, max_iou=1.0)
+            # replace NaNs (if any) with a large cost so Hungarian will avoid them
+            cost_for_assignment = np.where(np.isnan(cost_mat), 1e6, cost_mat)
+
+            # Hungarian assignment (minimize cost -> maximize IoU)
+            row_idx, col_idx = linear_sum_assignment(cost_for_assignment)
+
+            matched_gt = set()
+            matched_pred = set()
+            for r, c in zip(row_idx, col_idx):
+                if r < cost_mat.shape[0] and c < cost_mat.shape[1]:
+                    # IoU = 1 - cost
+                    cost_val = cost_mat[r, c]
+                    if np.isnan(cost_val):
+                        continue
+                    iou_val = 1.0 - float(cost_val)
+                    if iou_val >= self.iou_threshold:
+                        matched_gt.add(r)
+                        matched_pred.add(c)
+                        total_iou += iou_val
+                        iou_count += 1
+
+            frame_tp = len(matched_gt)
+            frame_fp = len(preds) - len(matched_pred)
+            frame_fn = len(gts) - len(matched_gt)
+
+            tp += frame_tp
+            fp += frame_fp
+            fn += frame_fn
+
+        precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
+        recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0
+        f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
+        avg_iou = total_iou / iou_count if iou_count > 0 else 0.0
+
+        total_gt = sum(len(v) for v in gt_tracks.values())
+        total_pred = sum(len(v) for v in pred_tracks.values())
+
+        return {
+            "precision": precision,
+            "recall": recall,
+            "f1": f1,
+            "avg_iou": avg_iou,
+            "tp": tp,
+            "fp": fp,
+            "fn": fn,
+            "total_gt": total_gt,
+            "total_pred": total_pred,
+        }
+
+    def _compute_mot_metrics(self, gt_tracks: Dict[int, List[Dict]], pred_tracks: Dict[int, List[Dict]]):
+        """
+        Use motmetrics.MOTAccumulator to collect associations per frame and compute common MOT metrics.
+        Distance matrix is taken directly from motmetrics.distances.iou_matrix (which returns 1 - IoU).
+        Pairs with distance > (1 - iou_threshold) are set to infinity to exclude them from matching.
+        """
+        acc = mm.MOTAccumulator(auto_id=True)
+
+        frames = sorted(set(list(gt_tracks.keys()) + list(pred_tracks.keys())))
+        for f in frames:
+            gts = gt_tracks.get(f, [])
+            preds = pred_tracks.get(f, [])
+
+            gt_ids = [g["track_id"] for g in gts]
+            pred_ids = [p["track_id"] for p in preds]
+
+            if gts and preds:
+                gt_boxes = np.array([g["bbox"] for g in gts])
+                pred_boxes = np.array([p["bbox"] for p in preds])
+
+                # motmetrics provides a distance matrix (1 - IoU)
+                dist_mat = mm.distances.iou_matrix(gt_boxes, pred_boxes, max_iou=1.0)
+                # exclude pairs with IoU < threshold => distance > 1 - threshold
+                dist_mat = np.array(dist_mat, dtype=float)
+                dist_mat[np.isnan(dist_mat)] = np.inf
+                dist_mat[dist_mat > (1.0 - self.iou_threshold)] = np.inf
+            else:
+                dist_mat = np.full((len(gts), len(preds)), np.inf)
+
+            acc.update(gt_ids, pred_ids, dist_mat)
+
+        mh = mm.metrics.create()
+        summary = mh.compute(
+            acc,
+            metrics=[
+                "mota",
+                "motp",
+                "idf1",
+                "num_switches",
+                "num_fragmentations",
+                "num_misses",
+                "num_false_positives",
+            ],
+            name="eval",
+        )
+
+        def get_val(col: str, default=0.0):
+            if summary.empty or col not in summary.columns:
+                return float(default)
+            v = summary.iloc[0][col]
+            return float(v) if not np.isnan(v) else float(default)
+
+        return {
+            "mota": get_val("mota", 0.0),
+            "motp": get_val("motp", 0.0),
+            "idf1": get_val("idf1", 0.0),
+            "id_switches": int(get_val("num_switches", 0.0)),
+            "fragmentations": int(get_val("num_fragmentations", 0.0)),
+            "num_misses": int(get_val("num_misses", 0.0)),
+            "num_false_positives": int(get_val("num_false_positives", 0.0)),
+        }
+
+
+def evaluate(
+    gt_annotation: VideoAnnotation,
+    pred_annotation: VideoAnnotation,
+    iou_threshold: float = 0.5,
+) -> Dict[str, Union[float, int]]:
+    """
+    Evaluate tracking predictions against ground truth.
+
+    Args:
+        gt_annotation: Ground-truth annotation, an object of class supervisely VideoAnnotation containing reference object tracks.
+        pred_annotation: Predicted annotation, an object of class supervisely VideoAnnotation to be compared against the ground truth.
+        iou_threshold: Minimum Intersection-over-Union required for a detection to be considered a valid match.
+
+    Returns:
+        dict: json with evaluation metrics.
+    """
+    evaluator = TrackingEvaluator(iou_threshold=iou_threshold)
+    return evaluator.evaluate(gt_annotation, pred_annotation)