arcsort 0.1.0__py3-none-any.whl

arcsort/cmc.py

@@ -0,0 +1,169 @@
+"""Camera Motion Compensation via sparse optical flow.
+
+Estimates the inter-frame camera motion as a 2D affine transform using
+Lucas-Kanade sparse optical flow and RANSAC-based affine estimation.
+The resulting transform is applied to Kalman filter predictions to
+compensate for camera panning, tilting, and zooming.
+
+This module is inspired by BoT-SORT's GMC (Global Motion Compensation)
+but simplified to only support sparse optical flow, which provides the
+best speed-accuracy trade-off for real-time tracking.
+"""
+
+from __future__ import annotations
+
+import cv2
+import numpy as np
+from loguru import logger
+
+
+# Default CMC parameters (overridden by config)
+DEFAULT_CMC_CONFIG = {
+    "max_corners": 1000,
+    "quality_level": 0.01,
+    "min_distance": 1,
+    "block_size": 3,
+    "ransac_threshold": 3.0,
+    "min_inliers": 10,
+    "downscale": 2,
+}
+
+
+class CameraMotionCompensator:
+    """Estimate inter-frame camera motion using sparse optical flow.
+
+    Tracks corner features between consecutive frames to estimate a
+    2D similarity transform (rotation, uniform scale, translation).
+    Falls back to identity transform when too few correspondences are
+    found (e.g., uniform backgrounds, heavy motion blur).
+    """
+
+    def __init__(self, config: dict | None = None) -> None:
+        """Initialize CMC module.
+
+        Args:
+            config: CMC configuration dictionary. Uses DEFAULT_CMC_CONFIG
+                for any missing keys.
+        """
+        if config is None:
+            config = {}
+        cfg = {**DEFAULT_CMC_CONFIG, **config}
+
+        self.downscale = max(1, int(cfg["downscale"]))
+        self.min_inliers = cfg["min_inliers"]
+        self.ransac_threshold = cfg["ransac_threshold"]
+
+        self.feature_params = {
+            "maxCorners": cfg["max_corners"],
+            "qualityLevel": cfg["quality_level"],
+            "minDistance": cfg["min_distance"],
+            "blockSize": cfg["block_size"],
+            "useHarrisDetector": False,
+            "k": 0.04,
+        }
+
+        self._prev_frame: np.ndarray | None = None
+        self._prev_keypoints: np.ndarray | None = None
+
+    def reset(self) -> None:
+        """Reset CMC state for a new sequence."""
+        self._prev_frame = None
+        self._prev_keypoints = None
+
+    def apply(self, frame: np.ndarray) -> np.ndarray:
+        """Compute affine transform between previous and current frame.
+
+        Args:
+            frame: Current frame as BGR image, shape (H, W, 3).
+
+        Returns:
+            2x3 affine matrix representing camera motion from previous
+            frame to current frame. Returns identity if this is the
+            first frame or if estimation fails.
+        """
+        height, width = frame.shape[:2]
+        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+        if self.downscale > 1:
+            gray = cv2.resize(
+                gray, (width // self.downscale, height // self.downscale)
+            )
+
+        keypoints = cv2.goodFeaturesToTrack(gray, mask=None, **self.feature_params)
+
+        # First frame: store and return identity
+        if self._prev_frame is None:
+            self._prev_frame = gray
+            self._prev_keypoints = keypoints
+            return np.eye(2, 3, dtype=np.float64)
+
+        affine = self._estimate_affine(gray, keypoints)
+
+        # Store for next frame
+        self._prev_frame = gray
+        self._prev_keypoints = keypoints
+
+        return affine
+
+    def _estimate_affine(
+        self, curr_frame: np.ndarray, curr_keypoints: np.ndarray | None
+    ) -> np.ndarray:
+        """Estimate affine transform via optical flow and RANSAC.
+
+        Args:
+            curr_frame: Current grayscale frame (possibly downscaled).
+            curr_keypoints: Feature points detected in current frame
+                (used for next iteration, not directly in flow).
+
+        Returns:
+            2x3 affine matrix. Identity if estimation fails.
+        """
+        identity = np.eye(2, 3, dtype=np.float64)
+
+        if self._prev_keypoints is None or len(self._prev_keypoints) < 4:
+            return identity
+
+        # Track previous keypoints into current frame
+        matched_pts, status, _err = cv2.calcOpticalFlowPyrLK(
+            self._prev_frame, curr_frame, self._prev_keypoints, None
+        )
+
+        if matched_pts is None:
+            return identity
+
+        # Keep only successfully tracked points
+        mask = status.flatten() == 1
+        prev_good = self._prev_keypoints[mask]
+        curr_good = matched_pts[mask]
+
+        if len(prev_good) < self.min_inliers:
+            logger.debug(
+                f"CMC: too few correspondences ({len(prev_good)}), "
+                f"using identity"
+            )
+            return identity
+
+        # Estimate similarity transform (4 DOF: rotation, scale, tx, ty)
+        affine, inliers = cv2.estimateAffinePartial2D(
+            prev_good, curr_good,
+            method=cv2.RANSAC,
+            ransacReprojThreshold=self.ransac_threshold,
+        )
+
+        if affine is None:
+            logger.debug("CMC: RANSAC failed, using identity")
+            return identity
+
+        num_inliers = int(inliers.sum()) if inliers is not None else 0
+        if num_inliers < self.min_inliers:
+            logger.debug(
+                f"CMC: too few inliers ({num_inliers}), using identity"
+            )
+            return identity
+
+        # Correct translation for downscaling
+        if self.downscale > 1:
+            affine[0, 2] *= self.downscale
+            affine[1, 2] *= self.downscale
+
+        return affine

arcsort/configs/__init__.py

@@ -0,0 +1,35 @@
+"""Tracker presets: the exact configurations behind the paper's results.
+
+Each preset is a YAML file shipped with the package:
+
+- ``sportsmot``: SportsMOT test row (81.82 HOTA); the recommended default for
+  sports broadcast footage.
+- ``dancetrack``: DanceTrack test row (62.1 HOTA).
+- ``sportsmot_deepeiou_inputs``: SportsMOT row on Deep-EIoU's released
+  detections and embeddings (80.31 HOTA); detector-calibration thresholds differ.
+"""
+
+from __future__ import annotations
+
+from importlib import resources
+
+import yaml
+
+
+def available() -> list[str]:
+    """Names of the shipped presets."""
+    files = resources.files(__name__).iterdir()
+    return sorted(f.name.removesuffix(".yaml") for f in files if f.name.endswith(".yaml"))
+
+
+def load_preset(name: str) -> dict:
+    """Load a preset's tracker config dict by name.
+
+    Raises:
+        ValueError: If no preset with that name exists.
+    """
+    path = resources.files(__name__) / f"{name}.yaml"
+    if not path.is_file():
+        raise ValueError(f"Unknown preset '{name}'. Available: {available()}")
+    config = yaml.safe_load(path.read_text())
+    return config.get("tracker", config)

arcsort/configs/dancetrack.yaml

@@ -0,0 +1,74 @@
+# ARC-SORT preset for DanceTrack (paper: DanceTrack test 62.1 HOTA).
+# Same association settings as sportsmot.yaml except the EIoU expansion
+# scale (1.0) and the FastReID SBS-S50 appearance model (see README: Weights).
+tracker:
+  name: ARC-SORT DanceTrack
+  track_high_thresh: 0.6
+  track_low_thresh: 0.1
+  new_track_thresh: 0.7
+  track_buffer: 30
+  match_thresh: 0.25
+  byte_match_thresh: 0.5
+  min_hits: 3
+  use_byte: true
+  use_cmc: true
+  use_imm: true
+  use_hybrid_distance: false
+  use_reid: true
+  use_gated_reid: false
+  reid:
+    weights: weights/dance_sbs_S50_traced.pt
+    crop_size:
+    - 128
+    - 384
+    w_assoc_emb: 1.25
+    alpha_fixed_emb: 0.95
+    aw_off: false
+    aw_param: 1.0
+    imm_gate_cv_thresh: 0.7
+    imm_gate_young_frames: 5
+    uncertainty_weight_scale: 1.0
+  imm:
+    models:
+    - cv
+    - ca
+    - ct
+    transition_prob_diagonal: 0.98
+    process_noise_cv: 1.0
+    process_noise_ca: 1.0
+    process_noise_ct: 1.0
+    initial_model_probs:
+    - 0.7
+    - 0.15
+    - 0.15
+  hybrid_distance:
+    iou_fallback_thresh: 0.1
+    shape_penalty_weight: 0.3
+  cmc:
+    max_corners: 1000
+    quality_level: 0.01
+    min_distance: 1
+    block_size: 3
+    ransac_threshold: 3.0
+    min_inliers: 10
+    downscale: 2
+  ocsort:
+    use_oru: true
+    use_ocm: true
+    use_ocr: true
+    delta_t: 3
+    inertia_weight: 0.2
+  use_eiou: true
+  eiou:
+    adaptive: true
+    adaptive_scale: 1.0
+    min_expansion: 0.2
+    max_expansion: 1.2
+  ambiguity:
+    enabled: true
+    margin_thresh: 0.2
+  arbitration:
+    enabled: true
+    window: 30
+    min_frames: 4
+    llr_thresh: 0.5

arcsort/configs/sportsmot.yaml

@@ -0,0 +1,80 @@
+# ARC-SORT preset for sports broadcast footage (paper: SportsMOT test 81.82 HOTA).
+# SportsMOT is 25 FPS, 720p footage with fast non-linear athlete motion and a
+# moving camera. All values were chosen on the validation split only.
+
+tracker:
+  name: "ARC-SORT SportsMOT"
+
+  track_high_thresh: 0.6
+  track_low_thresh: 0.1
+  new_track_thresh: 0.7
+  track_buffer: 30
+  match_thresh: 0.25
+  byte_match_thresh: 0.5
+  min_hits: 3
+
+  use_byte: true
+  use_cmc: true
+  use_imm: true
+  # Uncertainty-adaptive Expansion IoU is the stage-1 spatial metric
+  # (on val it beat both the hybrid distance and fixed-expansion EIoU)
+  use_hybrid_distance: false
+  use_eiou: true
+  use_reid: true
+  use_gated_reid: false
+  use_mahalanobis_gate: false
+
+  eiou:
+    adaptive: true
+    adaptive_scale: 2.0
+    min_expansion: 0.2
+    max_expansion: 1.2
+
+  reid:
+    # Sports-domain OSNet from the Deep-EIoU release, converted with
+    # tools/convert_osnet_weights.py (see README: Weights)
+    weights: "weights/sportsmot/sports_osnet_traced.pt"
+    crop_size: [128, 256]
+    w_assoc_emb: 1.25
+    alpha_fixed_emb: 0.95
+    aw_off: false
+    aw_param: 1.0
+    imm_gate_cv_thresh: 0.7
+    imm_gate_young_frames: 5
+
+  imm:
+    models: ["cv", "ca", "ct"]
+    transition_prob_diagonal: 0.98
+    process_noise_cv: 1.0
+    process_noise_ca: 1.0
+    process_noise_ct: 1.0
+    initial_model_probs: [0.7, 0.15, 0.15]
+
+  hybrid_distance:
+    iou_fallback_thresh: 0.1
+    shape_penalty_weight: 0.3
+
+  cmc:
+    max_corners: 1000
+    quality_level: 0.01
+    min_distance: 1
+    block_size: 3
+    ransac_threshold: 3.0
+    min_inliers: 10
+    downscale: 2
+
+  ocsort:
+    use_oru: true
+    use_ocm: true
+    use_ocr: true
+    use_cmc_oru: true
+    delta_t: 3
+    inertia_weight: 0.2
+  ambiguity:
+    enabled: true
+    margin_thresh: 0.2
+  arbitration:
+    enabled: true
+    window: 30
+    min_frames: 4
+    llr_thresh: 0.5

arcsort/configs/sportsmot_deepeiou_inputs.yaml

@@ -0,0 +1,68 @@
+# ARC-SORT on Deep-EIoU's released SportsMOT detections and embeddings
+# (paper: 80.31 HOTA, +3.1 over Deep-EIoU on identical inputs).
+# Same as sportsmot.yaml except the detector-calibration thresholds, which
+# match Deep-EIoU's detector: track_high 0.5, new_track 0.6, match 0.2.
+tracker:
+  name: ARC-SORT SportsMOT
+  track_high_thresh: 0.5
+  track_low_thresh: 0.1
+  new_track_thresh: 0.6
+  track_buffer: 30
+  match_thresh: 0.2
+  byte_match_thresh: 0.5
+  min_hits: 3
+  use_byte: true
+  use_cmc: true
+  use_imm: true
+  use_hybrid_distance: false
+  use_eiou: true
+  use_reid: true
+  use_gated_reid: false
+  use_mahalanobis_gate: false
+  eiou:
+    adaptive: true
+    adaptive_scale: 2.0
+    min_expansion: 0.2
+    max_expansion: 1.2
+  reid:
+    weights: weights/sportsmot/sports_osnet_traced.pt
+    crop_size:
+    - 128
+    - 256
+    w_assoc_emb: 1.25
+    alpha_fixed_emb: 0.95
+    aw_off: false
+    aw_param: 1.0
+    imm_gate_cv_thresh: 0.7
+    imm_gate_young_frames: 5
+  imm:
+    models:
+    - cv
+    - ca
+    - ct
+    transition_prob_diagonal: 0.98
+    process_noise_cv: 1.0
+    process_noise_ca: 1.0
+    process_noise_ct: 1.0
+    initial_model_probs:
+    - 0.7
+    - 0.15
+    - 0.15
+  hybrid_distance:
+    iou_fallback_thresh: 0.1
+    shape_penalty_weight: 0.3
+  cmc:
+    max_corners: 1000
+    quality_level: 0.01
+    min_distance: 1
+    block_size: 3
+    ransac_threshold: 3.0
+    min_inliers: 10
+    downscale: 2
+  ocsort:
+    use_oru: true
+    use_ocm: true
+    use_ocr: true
+    use_cmc_oru: true
+    delta_t: 3
+    inertia_weight: 0.2