matrice-analytics 0.1.60__py3-none-any.whl

matrice_analytics/post_processing/advanced_tracker/matching.py

@@ -0,0 +1,195 @@
+"""
+Matching utilities for track association.
+
+This module provides utilities for matching tracks with detections,
+including IoU distance calculation and linear assignment.
+"""
+
+import numpy as np
+import scipy
+from scipy.spatial.distance import cdist
+
+try:
+    import lap  # for linear_assignment
+    assert lap.__version__  # verify package is not directory
+except (ImportError, AssertionError, AttributeError):
+    # Fallback to scipy if lap is not available
+    lap = None
+
+
+def linear_assignment(cost_matrix: np.ndarray, thresh: float, use_lap: bool = True):
+    """
+    Perform linear assignment using either the scipy or lap.lapjv method.
+
+    Args:
+        cost_matrix (np.ndarray): The matrix containing cost values for assignments, with shape (N, M).
+        thresh (float): Threshold for considering an assignment valid.
+        use_lap (bool): Use lap.lapjv for the assignment. If False, scipy.optimize.linear_sum_assignment is used.
+
+    Returns:
+        matched_indices (np.ndarray): Array of matched indices of shape (K, 2), where K is the number of matches.
+        unmatched_a (np.ndarray): Array of unmatched indices from the first set, with shape (L,).
+        unmatched_b (np.ndarray): Array of unmatched indices from the second set, with shape (M,).
+    """
+    if cost_matrix.size == 0:
+        return np.empty((0, 2), dtype=int), tuple(range(cost_matrix.shape[0])), tuple(range(cost_matrix.shape[1]))
+
+    if use_lap and lap is not None:
+        # Use lap.lapjv
+        _, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh)
+        matches = [[ix, mx] for ix, mx in enumerate(x) if mx >= 0]
+        unmatched_a = np.where(x < 0)[0]
+        unmatched_b = np.where(y < 0)[0]
+    else:
+        # Use scipy.optimize.linear_sum_assignment
+        x, y = scipy.optimize.linear_sum_assignment(cost_matrix)  # row x, col y
+        matches = np.asarray([[x[i], y[i]] for i in range(len(x)) if cost_matrix[x[i], y[i]] <= thresh])
+        if len(matches) == 0:
+            unmatched_a = list(np.arange(cost_matrix.shape[0]))
+            unmatched_b = list(np.arange(cost_matrix.shape[1]))
+        else:
+            unmatched_a = list(frozenset(np.arange(cost_matrix.shape[0])) - frozenset(matches[:, 0]))
+            unmatched_b = list(frozenset(np.arange(cost_matrix.shape[1])) - frozenset(matches[:, 1]))
+
+    return matches, unmatched_a, unmatched_b
+
+
+def bbox_ioa(box1: np.ndarray, box2: np.ndarray, iou: bool = True) -> np.ndarray:
+    """
+    Calculate the intersection over area of box1, box2. Boxes are in x1y1x2y2 format.
+
+    Args:
+        box1 (np.ndarray): First set of boxes (N, 4)
+        box2 (np.ndarray): Second set of boxes (M, 4)
+        iou (bool): If True, calculate IoU, otherwise calculate IoA
+
+    Returns:
+        np.ndarray: IoU/IoA matrix of shape (N, M)
+    """
+    # Returns the intersection over box1 area by default
+    # box1: (N, 4), box2: (M, 4)
+    N = box1.shape[0]
+    M = box2.shape[0]
+
+    # Calculate intersection
+    tl = np.maximum(box1[:, None, :2], box2[:, :2])  # (N, M, 2)
+    br = np.minimum(box1[:, None, 2:], box2[:, 2:])  # (N, M, 2)
+    wh = np.maximum(0, br - tl)  # (N, M, 2)
+    inter = wh[:, :, 0] * wh[:, :, 1]  # (N, M)
+
+    # Calculate areas
+    area1 = (box1[:, 2] - box1[:, 0]) * (box1[:, 3] - box1[:, 1])  # (N,)
+    area2 = (box2[:, 2] - box2[:, 0]) * (box2[:, 3] - box2[:, 1])  # (M,)
+
+    if iou:
+        union = area1[:, None] + area2 - inter
+        return inter / union
+    else:
+        return inter / area1[:, None]
+
+
+def iou_distance(atracks: list, btracks: list) -> np.ndarray:
+    """
+    Compute cost based on Intersection over Union (IoU) between tracks.
+
+    Args:
+        atracks (List[STrack] or List[np.ndarray]): List of tracks 'a' or bounding boxes.
+        btracks (List[STrack] or List[np.ndarray]): List of tracks 'b' or bounding boxes.
+
+    Returns:
+        (np.ndarray): Cost matrix computed based on IoU with shape (len(atracks), len(btracks)).
+    """
+    if atracks and isinstance(atracks[0], np.ndarray) or btracks and isinstance(btracks[0], np.ndarray):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        # Extract bounding boxes from track objects
+        atlbrs = []
+        btlbrs = []
+        for track in atracks:
+            if hasattr(track, 'xyxy'):
+                atlbrs.append(track.xyxy)
+            elif hasattr(track, 'tlwh'):
+                tlwh = track.tlwh
+                atlbrs.append([tlwh[0], tlwh[1], tlwh[0] + tlwh[2], tlwh[1] + tlwh[3]])
+            else:
+                atlbrs.append([0, 0, 0, 0])
+        
+        for track in btracks:
+            if hasattr(track, 'xyxy'):
+                btlbrs.append(track.xyxy)
+            elif hasattr(track, 'tlwh'):
+                tlwh = track.tlwh
+                btlbrs.append([tlwh[0], tlwh[1], tlwh[0] + tlwh[2], tlwh[1] + tlwh[3]])
+            else:
+                btlbrs.append([0, 0, 0, 0])
+
+    ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+    if len(atlbrs) and len(btlbrs):
+        ious = bbox_ioa(
+            np.ascontiguousarray(atlbrs, dtype=np.float32),
+            np.ascontiguousarray(btlbrs, dtype=np.float32),
+            iou=True,
+        )
+    return 1 - ious  # cost matrix
+
+
+def embedding_distance(tracks: list, detections: list, metric: str = "cosine") -> np.ndarray:
+    """
+    Compute distance between tracks and detections based on embeddings.
+
+    Args:
+        tracks (List[STrack] or List[np.ndarray]): List of tracks, where each track contains embedding features.
+        detections (List[BaseTrack]): List of detections, where each detection contains embedding features.
+        metric (str): Metric for distance computation. Supported metrics include 'cosine', 'euclidean', etc.
+
+    Returns:
+        (np.ndarray): Cost matrix computed based on embeddings with shape (N, M), where N is the number of tracks
+            and M is the number of detections.
+    """
+    cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float32)
+    if cost_matrix.size == 0:
+        return cost_matrix
+    
+    # Extract features from detections
+    det_features = []
+    for det in detections:
+        if hasattr(det, 'curr_feat') and det.curr_feat is not None:
+            det_features.append(det.curr_feat)
+        else:
+            det_features.append(np.zeros(128))  # Default feature size
+    
+    det_features = np.asarray(det_features, dtype=np.float32)
+    
+    # Extract features from tracks
+    track_features = []
+    for track in tracks:
+        if hasattr(track, 'smooth_feat') and track.smooth_feat is not None:
+            track_features.append(track.smooth_feat)
+        else:
+            track_features.append(np.zeros(128))  # Default feature size
+    
+    track_features = np.asarray(track_features, dtype=np.float32)
+    
+    cost_matrix = np.maximum(0.0, cdist(track_features, det_features, metric))
+    return cost_matrix
+
+
+def fuse_score(cost_matrix: np.ndarray, detections: list) -> np.ndarray:
+    """
+    Fuse cost matrix with detection scores to produce a single similarity matrix.
+
+    Args:
+        cost_matrix (np.ndarray): The matrix containing cost values for assignments, with shape (N, M).
+        detections (List[BaseTrack]): List of detections, each containing a score attribute.
+
+    Returns:
+        (np.ndarray): Fused similarity matrix with shape (N, M).
+    """
+    if cost_matrix.size == 0:
+        return cost_matrix
+    iou_sim = 1 - cost_matrix
+    det_scores = np.array([det.score for det in detections])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    fuse_sim = iou_sim * det_scores
+    return 1 - fuse_sim  # fuse_cost 

matrice_analytics/post_processing/advanced_tracker/strack.py

@@ -0,0 +1,230 @@
+"""
+STrack class for single object tracking.
+
+This module provides the STrack class that represents a single tracked object
+with Kalman filtering for state estimation.
+"""
+
+from typing import Any, List
+import numpy as np
+
+from .base import BaseTrack, TrackState
+from .kalman_filter import KalmanFilterXYAH
+
+
+def xywh2ltwh(xywh: List[float]) -> List[float]:
+    """
+    Convert bounding box from center format (x, y, w, h) to top-left format (x, y, w, h).
+    
+    Args:
+        xywh (List[float]): Bounding box in center format [x, y, w, h]
+        
+    Returns:
+        List[float]: Bounding box in top-left format [x, y, w, h]
+    """
+    x, y, w, h = xywh
+    return [x - w/2, y - h/2, w, h]
+
+
+class STrack(BaseTrack):
+    """
+    Single object tracking representation that uses Kalman filtering for state estimation.
+
+    This class is responsible for storing all the information regarding individual tracklets and performs state updates
+    and predictions based on Kalman filter.
+
+    Attributes:
+        shared_kalman (KalmanFilterXYAH): Shared Kalman filter used across all STrack instances for prediction.
+        _tlwh (np.ndarray): Private attribute to store top-left corner coordinates and width and height of bounding box.
+        kalman_filter (KalmanFilterXYAH): Instance of Kalman filter used for this particular object track.
+        mean (np.ndarray): Mean state estimate vector.
+        covariance (np.ndarray): Covariance of state estimate.
+        is_activated (bool): Boolean flag indicating if the track has been activated.
+        score (float): Confidence score of the track.
+        tracklet_len (int): Length of the tracklet.
+        cls (Any): Class label for the object.
+        idx (int): Index or identifier for the object.
+        frame_id (int): Current frame ID.
+        start_frame (int): Frame where the object was first detected.
+        angle (float or None): Optional angle information for oriented bounding boxes.
+    """
+
+    shared_kalman = KalmanFilterXYAH()
+
+    def __init__(self, xywh: List[float], score: float, cls: Any):
+        """
+        Initialize a new STrack instance.
+
+        Args:
+            xywh (List[float]): Bounding box coordinates and dimensions in the format (x, y, w, h, [a], idx), where
+                (x, y) is the center, (w, h) are width and height, [a] is optional aspect ratio, and idx is the id.
+            score (float): Confidence score of the detection.
+            cls (Any): Class label for the detected object.
+        """
+        super().__init__()
+        # xywh+idx or xywha+idx
+        assert len(xywh) in {5, 6}, f"expected 5 or 6 values but got {len(xywh)}"
+        self._tlwh = np.asarray(xywh2ltwh(xywh[:4]), dtype=np.float32)
+        self.kalman_filter = None
+        self.mean, self.covariance = None, None
+        self.is_activated = False
+
+        self.score = score
+        self.tracklet_len = 0
+        self.cls = cls
+        self.idx = xywh[-1]
+        self.angle = xywh[4] if len(xywh) == 6 else None
+
+    def predict(self):
+        """Predict the next state (mean and covariance) of the object using the Kalman filter."""
+        mean_state = self.mean.copy()
+        if self.state != TrackState.Tracked:
+            mean_state[7] = 0
+        self.mean, self.covariance = self.kalman_filter.predict(mean_state, self.covariance)
+
+    @staticmethod
+    def multi_predict(stracks: List["STrack"]):
+        """Perform multi-object predictive tracking using Kalman filter for the provided list of STrack instances."""
+        if len(stracks) <= 0:
+            return
+        multi_mean = np.asarray([st.mean.copy() for st in stracks])
+        multi_covariance = np.asarray([st.covariance for st in stracks])
+        for i, st in enumerate(stracks):
+            if st.state != TrackState.Tracked:
+                multi_mean[i][7] = 0
+        multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(multi_mean, multi_covariance)
+        for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+            stracks[i].mean = mean
+            stracks[i].covariance = cov
+
+    @staticmethod
+    def multi_gmc(stracks: List["STrack"], H: np.ndarray = np.eye(2, 3)):
+        """Update state tracks positions and covariances using a homography matrix for multiple tracks."""
+        if len(stracks) > 0:
+            multi_mean = np.asarray([st.mean.copy() for st in stracks])
+            multi_covariance = np.asarray([st.covariance for st in stracks])
+
+            R = H[:2, :2]
+            R8x8 = np.kron(np.eye(4, dtype=float), R)
+            t = H[:2, 2]
+
+            for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+                mean = R8x8.dot(mean)
+                mean[:2] += t
+                cov = R8x8.dot(cov).dot(R8x8.transpose())
+
+                stracks[i].mean = mean
+                stracks[i].covariance = cov
+
+    def activate(self, kalman_filter: KalmanFilterXYAH, frame_id: int):
+        """Activate a new tracklet using the provided Kalman filter and initialize its state and covariance."""
+        self.kalman_filter = kalman_filter
+        self.track_id = self.next_id()
+        self.mean, self.covariance = self.kalman_filter.initiate(self.convert_coords(self._tlwh))
+
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        if frame_id == 1:
+            self.is_activated = True
+        self.frame_id = frame_id
+        self.start_frame = frame_id
+
+    def re_activate(self, new_track: "STrack", frame_id: int, new_id: bool = False):
+        """Reactivate a previously lost track using new detection data and update its state and attributes."""
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.convert_coords(new_track.tlwh)
+        )
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        self.is_activated = True
+        self.frame_id = frame_id
+        if new_id:
+            self.track_id = self.next_id()
+        self.score = new_track.score
+        self.cls = new_track.cls
+        self.angle = new_track.angle
+        self.idx = new_track.idx
+        
+        # CRITICAL FIX: Preserve original detection data for face recognition fields
+        if hasattr(new_track, 'original_detection'):
+            self.original_detection = new_track.original_detection
+
+    def update(self, new_track: "STrack", frame_id: int):
+        """
+        Update the state of a matched track.
+
+        Args:
+            new_track (STrack): The new track containing updated information.
+            frame_id (int): The ID of the current frame.
+        """
+        self.frame_id = frame_id
+        self.tracklet_len += 1
+
+        new_tlwh = new_track.tlwh
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.convert_coords(new_tlwh)
+        )
+        self.state = TrackState.Tracked
+        self.is_activated = True
+
+        self.score = new_track.score
+        self.cls = new_track.cls
+        self.angle = new_track.angle
+        self.idx = new_track.idx
+        
+        # CRITICAL FIX: Preserve original detection data for face recognition fields
+        if hasattr(new_track, 'original_detection'):
+            self.original_detection = new_track.original_detection
+
+    def convert_coords(self, tlwh: np.ndarray) -> np.ndarray:
+        """Convert a bounding box's top-left-width-height format to its x-y-aspect-height equivalent."""
+        return self.tlwh_to_xyah(tlwh)
+
+    @property
+    def tlwh(self) -> np.ndarray:
+        """Get the bounding box in top-left-width-height format from the current state estimate."""
+        if self.mean is None:
+            return self._tlwh.copy()
+        ret = self.mean[:4].copy()
+        ret[2] *= ret[3]
+        ret[:2] -= ret[2:] / 2
+        return ret
+
+    @property
+    def xyxy(self) -> np.ndarray:
+        """Convert bounding box from (top left x, top left y, width, height) to (min x, min y, max x, max y) format."""
+        ret = self.tlwh.copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    @staticmethod
+    def tlwh_to_xyah(tlwh: np.ndarray) -> np.ndarray:
+        """Convert bounding box from tlwh format to center-x-center-y-aspect-height (xyah) format."""
+        ret = np.asarray(tlwh).copy()
+        ret[:2] += ret[2:] / 2
+        ret[2] /= ret[3]
+        return ret
+
+    @property
+    def xywh(self) -> np.ndarray:
+        """Get the current position of the bounding box in (center x, center y, width, height) format."""
+        ret = np.asarray(self.tlwh).copy()
+        ret[:2] += ret[2:] / 2
+        return ret
+
+    @property
+    def xywha(self) -> np.ndarray:
+        """Get position in (center x, center y, width, height, angle) format, warning if angle is missing."""
+        if self.angle is None:
+            return self.xywh
+        return np.concatenate([self.xywh, self.angle[None]])
+
+    @property
+    def result(self) -> List[float]:
+        """Get the current tracking results in the appropriate bounding box format."""
+        coords = self.xyxy if self.angle is None else self.xywha
+        return coords.tolist() + [self.track_id, self.score, self.cls, self.idx]
+
+    def __repr__(self) -> str:
+        """Return a string representation of the STrack object including start frame, end frame, and track ID."""
+        return f"OT_{self.track_id}_({self.start_frame}-{self.end_frame})"