Reference for ultralytics/trackers/utils/gmc.py

def apply(self, raw_frame: np.array, detections: list = None) -> np.array:
    """
    Apply object detection on a raw frame using the specified method.

    Args:
        raw_frame (np.ndarray): The raw frame to be processed, with shape (H, W, C).
        detections (List | None): List of detections to be used in the processing.

    Returns:
        (np.ndarray): Processed frame with applied object detection.

    Examples:
        >>> gmc = GMC(method="sparseOptFlow")
        >>> raw_frame = np.random.rand(480, 640, 3)
        >>> processed_frame = gmc.apply(raw_frame)
        >>> print(processed_frame.shape)
        (480, 640, 3)
    """
    if self.method in {"orb", "sift"}:
        return self.applyFeatures(raw_frame, detections)
    elif self.method == "ecc":
        return self.applyEcc(raw_frame)
    elif self.method == "sparseOptFlow":
        return self.applySparseOptFlow(raw_frame)
    else:
        return np.eye(2, 3)

def applyEcc(self, raw_frame: np.array) -> np.array:
    """
    Apply the ECC (Enhanced Correlation Coefficient) algorithm to a raw frame for motion compensation.

    Args:
        raw_frame (np.ndarray): The raw frame to be processed, with shape (H, W, C).

    Returns:
        (np.ndarray): The processed frame with the applied ECC transformation.

    Examples:
        >>> gmc = GMC(method="ecc")
        >>> processed_frame = gmc.applyEcc(np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]))
        >>> print(processed_frame)
        [[1. 0. 0.]
         [0. 1. 0.]]
    """
    height, width, _ = raw_frame.shape
    frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2GRAY)
    H = np.eye(2, 3, dtype=np.float32)

    # Downscale image
    if self.downscale > 1.0:
        frame = cv2.GaussianBlur(frame, (3, 3), 1.5)
        frame = cv2.resize(frame, (width // self.downscale, height // self.downscale))

    # Handle first frame
    if not self.initializedFirstFrame:
        # Initialize data
        self.prevFrame = frame.copy()

        # Initialization done
        self.initializedFirstFrame = True

        return H

    # Run the ECC algorithm. The results are stored in warp_matrix.
    # (cc, H) = cv2.findTransformECC(self.prevFrame, frame, H, self.warp_mode, self.criteria)
    try:
        (_, H) = cv2.findTransformECC(self.prevFrame, frame, H, self.warp_mode, self.criteria, None, 1)
    except Exception as e:
        LOGGER.warning(f"WARNING: find transform failed. Set warp as identity {e}")

    return H

def applyFeatures(self, raw_frame: np.array, detections: list = None) -> np.array:
    """
    Apply feature-based methods like ORB or SIFT to a raw frame.

    Args:
        raw_frame (np.ndarray): The raw frame to be processed, with shape (H, W, C).
        detections (List | None): List of detections to be used in the processing.

    Returns:
        (np.ndarray): Processed frame.

    Examples:
        >>> gmc = GMC(method="orb")
        >>> raw_frame = np.random.randint(0, 255, (480, 640, 3), dtype=np.uint8)
        >>> processed_frame = gmc.applyFeatures(raw_frame)
        >>> print(processed_frame.shape)
        (2, 3)
    """
    height, width, _ = raw_frame.shape
    frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2GRAY)
    H = np.eye(2, 3)

    # Downscale image
    if self.downscale > 1.0:
        frame = cv2.resize(frame, (width // self.downscale, height // self.downscale))
        width = width // self.downscale
        height = height // self.downscale

    # Find the keypoints
    mask = np.zeros_like(frame)
    mask[int(0.02 * height) : int(0.98 * height), int(0.02 * width) : int(0.98 * width)] = 255
    if detections is not None:
        for det in detections:
            tlbr = (det[:4] / self.downscale).astype(np.int_)
            mask[tlbr[1] : tlbr[3], tlbr[0] : tlbr[2]] = 0

    keypoints = self.detector.detect(frame, mask)

    # Compute the descriptors
    keypoints, descriptors = self.extractor.compute(frame, keypoints)

    # Handle first frame
    if not self.initializedFirstFrame:
        # Initialize data
        self.prevFrame = frame.copy()
        self.prevKeyPoints = copy.copy(keypoints)
        self.prevDescriptors = copy.copy(descriptors)

        # Initialization done
        self.initializedFirstFrame = True

        return H

    # Match descriptors
    knnMatches = self.matcher.knnMatch(self.prevDescriptors, descriptors, 2)

    # Filter matches based on smallest spatial distance
    matches = []
    spatialDistances = []

    maxSpatialDistance = 0.25 * np.array([width, height])

    # Handle empty matches case
    if len(knnMatches) == 0:
        # Store to next iteration
        self.prevFrame = frame.copy()
        self.prevKeyPoints = copy.copy(keypoints)
        self.prevDescriptors = copy.copy(descriptors)

        return H

    for m, n in knnMatches:
        if m.distance < 0.9 * n.distance:
            prevKeyPointLocation = self.prevKeyPoints[m.queryIdx].pt
            currKeyPointLocation = keypoints[m.trainIdx].pt

            spatialDistance = (
                prevKeyPointLocation[0] - currKeyPointLocation[0],
                prevKeyPointLocation[1] - currKeyPointLocation[1],
            )

            if (np.abs(spatialDistance[0]) < maxSpatialDistance[0]) and (
                np.abs(spatialDistance[1]) < maxSpatialDistance[1]
            ):
                spatialDistances.append(spatialDistance)
                matches.append(m)

    meanSpatialDistances = np.mean(spatialDistances, 0)
    stdSpatialDistances = np.std(spatialDistances, 0)

    inliers = (spatialDistances - meanSpatialDistances) < 2.5 * stdSpatialDistances

    goodMatches = []
    prevPoints = []
    currPoints = []
    for i in range(len(matches)):
        if inliers[i, 0] and inliers[i, 1]:
            goodMatches.append(matches[i])
            prevPoints.append(self.prevKeyPoints[matches[i].queryIdx].pt)
            currPoints.append(keypoints[matches[i].trainIdx].pt)

    prevPoints = np.array(prevPoints)
    currPoints = np.array(currPoints)

    # Draw the keypoint matches on the output image
    # if False:
    #     import matplotlib.pyplot as plt
    #     matches_img = np.hstack((self.prevFrame, frame))
    #     matches_img = cv2.cvtColor(matches_img, cv2.COLOR_GRAY2BGR)
    #     W = self.prevFrame.shape[1]
    #     for m in goodMatches:
    #         prev_pt = np.array(self.prevKeyPoints[m.queryIdx].pt, dtype=np.int_)
    #         curr_pt = np.array(keypoints[m.trainIdx].pt, dtype=np.int_)
    #         curr_pt[0] += W
    #         color = np.random.randint(0, 255, 3)
    #         color = (int(color[0]), int(color[1]), int(color[2]))
    #
    #         matches_img = cv2.line(matches_img, prev_pt, curr_pt, tuple(color), 1, cv2.LINE_AA)
    #         matches_img = cv2.circle(matches_img, prev_pt, 2, tuple(color), -1)
    #         matches_img = cv2.circle(matches_img, curr_pt, 2, tuple(color), -1)
    #
    #     plt.figure()
    #     plt.imshow(matches_img)
    #     plt.show()

    # Find rigid matrix
    if prevPoints.shape[0] > 4:
        H, inliers = cv2.estimateAffinePartial2D(prevPoints, currPoints, cv2.RANSAC)

        # Handle downscale
        if self.downscale > 1.0:
            H[0, 2] *= self.downscale
            H[1, 2] *= self.downscale
    else:
        LOGGER.warning("WARNING: not enough matching points")

    # Store to next iteration
    self.prevFrame = frame.copy()
    self.prevKeyPoints = copy.copy(keypoints)
    self.prevDescriptors = copy.copy(descriptors)

    return H

def applySparseOptFlow(self, raw_frame: np.array) -> np.array:
    """
    Apply Sparse Optical Flow method to a raw frame.

    Args:
        raw_frame (np.ndarray): The raw frame to be processed, with shape (H, W, C).

    Returns:
        (np.ndarray): Processed frame with shape (2, 3).

    Examples:
        >>> gmc = GMC()
        >>> result = gmc.applySparseOptFlow(np.array([[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]))
        >>> print(result)
        [[1. 0. 0.]
         [0. 1. 0.]]
    """
    height, width, _ = raw_frame.shape
    frame = cv2.cvtColor(raw_frame, cv2.COLOR_BGR2GRAY)
    H = np.eye(2, 3)

    # Downscale image
    if self.downscale > 1.0:
        frame = cv2.resize(frame, (width // self.downscale, height // self.downscale))

    # Find the keypoints
    keypoints = cv2.goodFeaturesToTrack(frame, mask=None, **self.feature_params)

    # Handle first frame
    if not self.initializedFirstFrame or self.prevKeyPoints is None:
        self.prevFrame = frame.copy()
        self.prevKeyPoints = copy.copy(keypoints)
        self.initializedFirstFrame = True
        return H

    # Find correspondences
    matchedKeypoints, status, _ = cv2.calcOpticalFlowPyrLK(self.prevFrame, frame, self.prevKeyPoints, None)

    # Leave good correspondences only
    prevPoints = []
    currPoints = []

    for i in range(len(status)):
        if status[i]:
            prevPoints.append(self.prevKeyPoints[i])
            currPoints.append(matchedKeypoints[i])

    prevPoints = np.array(prevPoints)
    currPoints = np.array(currPoints)

    # Find rigid matrix
    if (prevPoints.shape[0] > 4) and (prevPoints.shape[0] == prevPoints.shape[0]):
        H, _ = cv2.estimateAffinePartial2D(prevPoints, currPoints, cv2.RANSAC)

        if self.downscale > 1.0:
            H[0, 2] *= self.downscale
            H[1, 2] *= self.downscale
    else:
        LOGGER.warning("WARNING: not enough matching points")

    self.prevFrame = frame.copy()
    self.prevKeyPoints = copy.copy(keypoints)

    return H

def reset_params(self) -> None:
    """Reset the internal parameters including previous frame, keypoints, and descriptors."""
    self.prevFrame = None
    self.prevKeyPoints = None
    self.prevDescriptors = None
    self.initializedFirstFrame = False