eye-cv 1.0.0__py3-none-any.whl

eye/tracker/byte_tracker/kalman_filter.py

@@ -0,0 +1,205 @@
+from typing import Tuple
+
+import numpy as np
+import scipy.linalg
+
+
+class KalmanFilter:
+    """
+    A simple Kalman filter for tracking bounding boxes in image space.
+
+    The 8-dimensional state space
+
+        x, y, a, h, vx, vy, va, vh
+
+    contains the bounding box center position (x, y), aspect ratio a, height h,
+    and their respective velocities.
+
+    Object motion follows a constant velocity model. The bounding box location
+    (x, y, a, h) is taken as direct observation of the state space (linear
+    observation model).
+    """
+
+    def __init__(self):
+        ndim, dt = 4, 1.0
+
+        self._motion_mat = np.eye(2 * ndim, 2 * ndim)
+        for i in range(ndim):
+            self._motion_mat[i, ndim + i] = dt
+        self._update_mat = np.eye(ndim, 2 * ndim)
+        self._std_weight_position = 1.0 / 20
+        self._std_weight_velocity = 1.0 / 160
+
+    def initiate(self, measurement: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Create track from an unassociated measurement.
+
+        Args:
+            measurement (ndarray): Bounding box coordinates (x, y, a, h) with
+                center position (x, y), aspect ratio a, and height h.
+
+        Returns:
+            Tuple[ndarray, ndarray]: Returns the mean vector (8 dimensional) and
+                covariance matrix (8x8 dimensional) of the new track.
+                Unobserved velocities are initialized to 0 mean.
+        """
+        mean_pos = measurement
+        mean_vel = np.zeros_like(mean_pos)
+        mean = np.r_[mean_pos, mean_vel]
+
+        std = [
+            2 * self._std_weight_position * measurement[3],
+            2 * self._std_weight_position * measurement[3],
+            1e-2,
+            2 * self._std_weight_position * measurement[3],
+            10 * self._std_weight_velocity * measurement[3],
+            10 * self._std_weight_velocity * measurement[3],
+            1e-5,
+            10 * self._std_weight_velocity * measurement[3],
+        ]
+        covariance = np.diag(np.square(std))
+        return mean, covariance
+
+    def predict(
+        self, mean: np.ndarray, covariance: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Run Kalman filter prediction step.
+
+        Args:
+            mean (ndarray): The 8 dimensional mean vector of the object
+                state at the previous time step.
+            covariance (ndarray): The 8x8 dimensional covariance matrix of
+                the object state at the previous time step.
+
+        Returns:
+            Tuple[ndarray, ndarray]: Returns the mean vector and
+                covariance matrix of the predicted state.
+                Unobserved velocities are initialized to 0 mean.
+        """
+        std_pos = [
+            self._std_weight_position * mean[3],
+            self._std_weight_position * mean[3],
+            1e-2,
+            self._std_weight_position * mean[3],
+        ]
+        std_vel = [
+            self._std_weight_velocity * mean[3],
+            self._std_weight_velocity * mean[3],
+            1e-5,
+            self._std_weight_velocity * mean[3],
+        ]
+        motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
+
+        mean = np.dot(mean, self._motion_mat.T)
+        covariance = (
+            np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T))
+            + motion_cov
+        )
+
+        return mean, covariance
+
+    def project(
+        self, mean: np.ndarray, covariance: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Project state distribution to measurement space.
+
+        Args:
+            mean (ndarray): The state's mean vector (8 dimensional array).
+            covariance (ndarray): The state's covariance matrix (8x8 dimensional).
+
+        Returns:
+            Tuple[ndarray, ndarray]: Returns the projected mean and
+                covariance matrix of the given state estimate.
+        """
+        std = [
+            self._std_weight_position * mean[3],
+            self._std_weight_position * mean[3],
+            1e-1,
+            self._std_weight_position * mean[3],
+        ]
+        innovation_cov = np.diag(np.square(std))
+
+        mean = np.dot(self._update_mat, mean)
+        covariance = np.linalg.multi_dot(
+            (self._update_mat, covariance, self._update_mat.T)
+        )
+        return mean, covariance + innovation_cov
+
+    def multi_predict(
+        self, mean: np.ndarray, covariance: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Run Kalman filter prediction step (Vectorized version).
+
+        Args:
+            mean (ndarray): The Nx8 dimensional mean matrix
+                of the object states at the previous time step.
+            covariance (ndarray): The Nx8x8 dimensional covariance matrices
+                of the object states at the previous time step.
+
+        Returns:
+            Tuple[ndarray, ndarray]: Returns the mean vector and
+                covariance matrix of the predicted state.
+                Unobserved velocities are initialized to 0 mean.
+        """
+        std_pos = [
+            self._std_weight_position * mean[:, 3],
+            self._std_weight_position * mean[:, 3],
+            1e-2 * np.ones_like(mean[:, 3]),
+            self._std_weight_position * mean[:, 3],
+        ]
+        std_vel = [
+            self._std_weight_velocity * mean[:, 3],
+            self._std_weight_velocity * mean[:, 3],
+            1e-5 * np.ones_like(mean[:, 3]),
+            self._std_weight_velocity * mean[:, 3],
+        ]
+        sqr = np.square(np.r_[std_pos, std_vel]).T
+
+        motion_cov = []
+        for i in range(len(mean)):
+            motion_cov.append(np.diag(sqr[i]))
+        motion_cov = np.asarray(motion_cov)
+
+        mean = np.dot(mean, self._motion_mat.T)
+        left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
+        covariance = np.dot(left, self._motion_mat.T) + motion_cov
+
+        return mean, covariance
+
+    def update(
+        self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Run Kalman filter correction step.
+
+        Args:
+            mean (ndarray): The predicted state's mean vector (8 dimensional).
+            covariance (ndarray): The state's covariance matrix (8x8 dimensional).
+            measurement (ndarray): The 4-dimensional measurement vector (x, y, a, h),
+                where (x, y) is the center position, a the aspect ratio,
+                and h the height of the bounding box.
+
+        Returns:
+            Tuple[ndarray, ndarray]: Returns the measurement-corrected
+                state distribution.
+        """
+        projected_mean, projected_cov = self.project(mean, covariance)
+
+        chol_factor, lower = scipy.linalg.cho_factor(
+            projected_cov, lower=True, check_finite=False
+        )
+        kalman_gain = scipy.linalg.cho_solve(
+            (chol_factor, lower),
+            np.dot(covariance, self._update_mat.T).T,
+            check_finite=False,
+        ).T
+        innovation = measurement - projected_mean
+
+        new_mean = mean + np.dot(innovation, kalman_gain.T)
+        new_covariance = covariance - np.linalg.multi_dot(
+            (kalman_gain, projected_cov, kalman_gain.T)
+        )
+        return new_mean, new_covariance

eye/tracker/byte_tracker/matching.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Tuple
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+from eye.detection.utils import box_iou_batch
+
+if TYPE_CHECKING:
+    from eye.tracker.byte_tracker.core import STrack
+
+
+def indices_to_matches(
+    cost_matrix: np.ndarray, indices: np.ndarray, thresh: float
+) -> Tuple[np.ndarray, tuple, tuple]:
+    matched_cost = cost_matrix[tuple(zip(*indices))]
+    matched_mask = matched_cost <= thresh
+
+    matches = indices[matched_mask]
+    unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
+    unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
+    return matches, unmatched_a, unmatched_b
+
+
+def linear_assignment(
+    cost_matrix: np.ndarray, thresh: float
+) -> Tuple[np.ndarray, Tuple[int], Tuple[int, int]]:
+    if cost_matrix.size == 0:
+        return (
+            np.empty((0, 2), dtype=int),
+            tuple(range(cost_matrix.shape[0])),
+            tuple(range(cost_matrix.shape[1])),
+        )
+
+    cost_matrix[cost_matrix > thresh] = thresh + 1e-4
+    row_ind, col_ind = linear_sum_assignment(cost_matrix)
+    indices = np.column_stack((row_ind, col_ind))
+
+    return indices_to_matches(cost_matrix, indices, thresh)
+
+
+def iou_distance(atracks: List[STrack], btracks: List[STrack]) -> np.ndarray:
+    if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
+        len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
+    ):
+        atlbrs = atracks
+        btlbrs = btracks
+    else:
+        atlbrs = [track.tlbr for track in atracks]
+        btlbrs = [track.tlbr for track in btracks]
+
+    _ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
+    if _ious.size != 0:
+        _ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs))
+    cost_matrix = 1 - _ious
+
+    return cost_matrix
+
+
+def fuse_score(cost_matrix: np.ndarray, stracks: List[STrack]) -> np.ndarray:
+    if cost_matrix.size == 0:
+        return cost_matrix
+    iou_sim = 1 - cost_matrix
+    det_scores = np.array([strack.score for strack in stracks])
+    det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
+    fuse_sim = iou_sim * det_scores
+    fuse_cost = 1 - fuse_sim
+    return fuse_cost

eye/tracker/byte_tracker/single_object_track.py

@@ -0,0 +1,178 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import List
+
+import numpy as np
+import numpy.typing as npt
+
+from eye.tracker.byte_tracker.kalman_filter import KalmanFilter
+from eye.tracker.byte_tracker.utils import IdCounter
+
+
+class TrackState(Enum):
+    New = 0
+    Tracked = 1
+    Lost = 2
+    Removed = 3
+
+
+class STrack:
+    def __init__(
+        self,
+        tlwh: npt.NDArray[np.float32],
+        score: npt.NDArray[np.float32],
+        minimum_consecutive_frames: int,
+        shared_kalman: KalmanFilter,
+        internal_id_counter: IdCounter,
+        external_id_counter: IdCounter,
+    ):
+        self.state = TrackState.New
+        self.is_activated = False
+        self.start_frame = 0
+        self.frame_id = 0
+
+        self._tlwh = np.asarray(tlwh, dtype=np.float32)
+        self.kalman_filter = None
+        self.shared_kalman = shared_kalman
+        self.mean, self.covariance = None, None
+        self.is_activated = False
+
+        self.score = score
+        self.tracklet_len = 0
+
+        self.minimum_consecutive_frames = minimum_consecutive_frames
+
+        self.internal_id_counter = internal_id_counter
+        self.external_id_counter = external_id_counter
+        self.internal_track_id = self.internal_id_counter.NO_ID
+        self.external_track_id = self.external_id_counter.NO_ID
+
+    def predict(self) -> None:
+        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], shared_kalman: KalmanFilter) -> None:
+        if len(stracks) > 0:
+            multi_mean = []
+            multi_covariance = []
+            for i, st in enumerate(stracks):
+                multi_mean.append(st.mean.copy())
+                multi_covariance.append(st.covariance)
+                if st.state != TrackState.Tracked:
+                    multi_mean[i][7] = 0
+
+            multi_mean, multi_covariance = shared_kalman.multi_predict(
+                np.asarray(multi_mean), np.asarray(multi_covariance)
+            )
+            for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
+                stracks[i].mean = mean
+                stracks[i].covariance = cov
+
+    def activate(self, kalman_filter: KalmanFilter, frame_id: int) -> None:
+        """Start a new tracklet"""
+        self.kalman_filter = kalman_filter
+        self.internal_track_id = self.internal_id_counter.new_id()
+        self.mean, self.covariance = self.kalman_filter.initiate(
+            self.tlwh_to_xyah(self._tlwh)
+        )
+
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+        if frame_id == 1:
+            self.is_activated = True
+
+        if self.minimum_consecutive_frames == 1:
+            self.external_track_id = self.external_id_counter.new_id()
+
+        self.frame_id = frame_id
+        self.start_frame = frame_id
+
+    def re_activate(self, new_track: STrack, frame_id: int) -> None:
+        self.mean, self.covariance = self.kalman_filter.update(
+            self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
+        )
+        self.tracklet_len = 0
+        self.state = TrackState.Tracked
+
+        self.frame_id = frame_id
+        self.score = new_track.score
+
+    def update(self, new_track: STrack, frame_id: int) -> None:
+        """
+        Update a matched track
+        :type new_track: STrack
+        :type frame_id: int
+        :type update_feature: bool
+        :return:
+        """
+        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.tlwh_to_xyah(new_tlwh)
+        )
+        self.state = TrackState.Tracked
+        if self.tracklet_len == self.minimum_consecutive_frames:
+            self.is_activated = True
+            if self.external_track_id == self.external_id_counter.NO_ID:
+                self.external_track_id = self.external_id_counter.new_id()
+
+        self.score = new_track.score
+
+    @property
+    def tlwh(self) -> npt.NDArray[np.float32]:
+        """Get current position in bounding box format `(top left x, top left y,
+        width, height)`.
+        """
+        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 tlbr(self) -> npt.NDArray[np.float32]:
+        """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
+        `(top left, bottom right)`.
+        """
+        ret = self.tlwh.copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    @staticmethod
+    def tlwh_to_xyah(tlwh) -> npt.NDArray[np.float32]:
+        """Convert bounding box to format `(center x, center y, aspect ratio,
+        height)`, where the aspect ratio is `width / height`.
+        """
+        ret = np.asarray(tlwh).copy()
+        ret[:2] += ret[2:] / 2
+        ret[2] /= ret[3]
+        return ret
+
+    def to_xyah(self) -> npt.NDArray[np.float32]:
+        return self.tlwh_to_xyah(self.tlwh)
+
+    @staticmethod
+    def tlbr_to_tlwh(tlbr) -> npt.NDArray[np.float32]:
+        ret = np.asarray(tlbr).copy()
+        ret[2:] -= ret[:2]
+        return ret
+
+    @staticmethod
+    def tlwh_to_tlbr(tlwh) -> npt.NDArray[np.float32]:
+        ret = np.asarray(tlwh).copy()
+        ret[2:] += ret[:2]
+        return ret
+
+    def __repr__(self) -> str:
+        return "OT_{}_({}-{})".format(
+            self.internal_track_id, self.start_frame, self.frame_id
+        )

eye/tracker/byte_tracker/utils.py

@@ -0,0 +1,18 @@
+class IdCounter:
+    def __init__(self, start_id: int = 0):
+        self.start_id = start_id
+        if self.start_id <= self.NO_ID:
+            raise ValueError(f"start_id must be greater than {self.NO_ID}")
+        self.reset()
+
+    def reset(self) -> None:
+        self._id = self.start_id
+
+    def new_id(self) -> int:
+        returned_id = self._id
+        self._id += 1
+        return returned_id
+
+    @property
+    def NO_ID(self) -> int:
+        return -1

eye/utils/conversion.py

@@ -0,0 +1,132 @@
+from functools import wraps
+from typing import List
+
+import cv2
+import numpy as np
+from PIL import Image
+
+from eye.annotators.base import ImageType
+
+
+def ensure_cv2_image_for_annotation(annotate_func):
+    """
+    Decorates `BaseAnnotator.annotate` implementations, converts scene to
+    an image type used internally by the annotators, converts back when annotation
+    is complete.
+
+    Assumes the annotators modify the scene in-place.
+    """
+
+    @wraps(annotate_func)
+    def wrapper(self, scene: ImageType, *args, **kwargs):
+        if isinstance(scene, np.ndarray):
+            return annotate_func(self, scene, *args, **kwargs)
+
+        if isinstance(scene, Image.Image):
+            scene_np = pillow_to_cv2(scene)
+            annotated_np = annotate_func(self, scene_np, *args, **kwargs)
+            scene.paste(cv2_to_pillow(annotated_np))
+            return scene
+
+        raise ValueError(f"Unsupported image type: {type(scene)}")
+
+    return wrapper
+
+
+def ensure_cv2_image_for_processing(image_processing_fun):
+    """
+    Decorates image processing functions that accept np.ndarray, converting `image` to
+    np.ndarray, converts back when processing is complete.
+
+    Assumes the annotators do NOT modify the scene in-place.
+    """
+
+    @wraps(image_processing_fun)
+    def wrapper(image: ImageType, *args, **kwargs):
+        if isinstance(image, np.ndarray):
+            return image_processing_fun(image, *args, **kwargs)
+
+        if isinstance(image, Image.Image):
+            scene = pillow_to_cv2(image)
+            annotated = image_processing_fun(scene, *args, **kwargs)
+            return cv2_to_pillow(annotated)
+
+        raise ValueError(f"Unsupported image type: {type(image)}")
+
+    return wrapper
+
+
+def ensure_pil_image_for_annotation(annotate_func):
+    """
+    Decorates image processing functions that accept np.ndarray, converting `image` to
+    PIL image, converts back when processing is complete.
+
+    Assumes the annotators modify the scene in-place.
+    """
+
+    @wraps(annotate_func)
+    def wrapper(self, scene: ImageType, *args, **kwargs):
+        if isinstance(scene, np.ndarray):
+            scene_pil = cv2_to_pillow(scene)
+            annotated_pil = annotate_func(self, scene_pil, *args, **kwargs)
+            np.copyto(scene, pillow_to_cv2(annotated_pil))
+            return scene
+
+        if isinstance(scene, Image.Image):
+            return annotate_func(self, scene, *args, **kwargs)
+
+        raise ValueError(f"Unsupported image type: {type(scene)}")
+
+    return wrapper
+
+
+def images_to_cv2(images: List[ImageType]) -> List[np.ndarray]:
+    """
+    Converts images provided either as Pillow images or OpenCV
+    images into OpenCV format.
+
+    Args:
+        images (List[ImageType]): Images to be converted
+
+    Returns:
+        List[np.ndarray]: List of input images in OpenCV format
+            (with order preserved).
+
+    """
+    result = []
+    for image in images:
+        if issubclass(type(image), Image.Image):
+            image = pillow_to_cv2(image)
+        result.append(image)
+    return result
+
+
+def pillow_to_cv2(image: Image.Image) -> np.ndarray:
+    """
+    Converts Pillow image into OpenCV image, handling RGB -> BGR
+    conversion.
+
+    Args:
+        image (Image.Image): Pillow image (in RGB format).
+
+    Returns:
+        (np.ndarray): Input image converted to OpenCV format.
+    """
+    scene = np.array(image)
+    scene = cv2.cvtColor(scene, cv2.COLOR_RGB2BGR)
+    return scene
+
+
+def cv2_to_pillow(image: np.ndarray) -> Image.Image:
+    """
+    Converts OpenCV image into Pillow image, handling BGR -> RGB
+    conversion.
+
+    Args:
+        image (np.ndarray): OpenCV image (in BGR format).
+
+    Returns:
+        (Image.Image): Input image converted to Pillow format.
+    """
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(image)