Spaces:
Sleeping
Sleeping
| # Mikel Broström 🔥 Yolo Tracking 🧾 AGPL-3.0 license | |
| import time | |
| import cv2 | |
| import numpy as np | |
| from boxmot.motion.cmc.base_cmc import BaseCMC | |
| from boxmot.utils import BOXMOT | |
| from boxmot.utils import logger as LOGGER | |
| class ECC(BaseCMC): | |
| def __init__( | |
| self, | |
| warp_mode: int = cv2.MOTION_EUCLIDEAN, | |
| eps: float = 1e-5, | |
| max_iter: int = 100, | |
| scale: float = 0.1, | |
| align: bool = False, | |
| grayscale: bool = True | |
| ) -> None: | |
| """Compute the warp matrix from src to dst. | |
| Parameters | |
| ---------- | |
| warp_mode: opencv flag | |
| translation: cv2.MOTION_TRANSLATION | |
| rotated and shifted: cv2.MOTION_EUCLIDEAN | |
| affine(shift,rotated,shear): cv2.MOTION_AFFINE | |
| homography(3d): cv2.MOTION_HOMOGRAPHY | |
| eps: float | |
| the threshold of the increment in the correlation coefficient between two iterations | |
| max_iter: int | |
| the number of iterations. | |
| scale: float or [int, int] | |
| scale_ratio: float | |
| scale_size: [W, H] | |
| align: bool | |
| whether to warp affine or perspective transforms to the source image | |
| grayscale: bool | |
| whether to transform 3 channel RGB to single channel grayscale for faster computations | |
| Returns | |
| ------- | |
| warp matrix : ndarray | |
| Returns the warp matrix from src to dst. | |
| if motion models is homography, the warp matrix will be 3x3, otherwise 2x3 | |
| src_aligned: ndarray | |
| aligned source image of gray | |
| """ | |
| self.align = align | |
| self.grayscale = grayscale | |
| self.scale = scale | |
| self.warp_mode = warp_mode | |
| self.termination_criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, max_iter, eps) | |
| self.prev_img = None | |
| def apply(self, img: np.ndarray, dets: np.ndarray = None) -> np.ndarray: | |
| """Apply sparse optical flow to compute the warp matrix. | |
| Parameters: | |
| img (ndarray): The input image. | |
| dets: Description of dets parameter. | |
| Returns: | |
| ndarray: The warp matrix from the source to the destination. | |
| If the motion model is homography, the warp matrix will be 3x3; otherwise, it will be 2x3. | |
| """ | |
| if self.warp_mode == cv2.MOTION_HOMOGRAPHY: | |
| warp_matrix = np.eye(3, 3, dtype=np.float32) | |
| else: | |
| warp_matrix = np.eye(2, 3, dtype=np.float32) | |
| if self.prev_img is None: | |
| self.prev_img = self.preprocess(img) | |
| return warp_matrix | |
| img = self.preprocess(img) | |
| try: | |
| (ret_val, warp_matrix) = cv2.findTransformECC( | |
| self.prev_img, # already processed | |
| img, | |
| warp_matrix, | |
| self.warp_mode, | |
| self.termination_criteria, | |
| None, | |
| 1 | |
| ) | |
| except Exception as e: | |
| LOGGER.warning(f'Affine matrix could not be generated: {e}. Returning identity') | |
| return warp_matrix | |
| # upscale warp matrix to original images size | |
| if self.scale < 1: | |
| warp_matrix[0, 2] /= self.scale | |
| warp_matrix[1, 2] /= self.scale | |
| if self.align: | |
| h, w = self.prev_img.shape | |
| if self.warp_mode == cv2.MOTION_HOMOGRAPHY: | |
| # Use warpPerspective for Homography | |
| self.prev_img_aligned = cv2.warpPerspective(self.prev_img, warp_matrix, (w, h), flags=cv2.INTER_LINEAR) | |
| else: | |
| # Use warpAffine for Translation, Euclidean and Affine | |
| self.prev_img_aligned = cv2.warpAffine(self.prev_img, warp_matrix, (w, h), flags=cv2.INTER_LINEAR) | |
| else: | |
| self.prev_img_aligned = None | |
| self.prev_img = img | |
| return warp_matrix # , prev_img_aligned | |
| def main(): | |
| ecc = ECC(scale=0.5, align=True, grayscale=True) | |
| curr_img = cv2.imread('assets/MOT17-mini/train/MOT17-2-FRCNN/img1/000005.jpg') | |
| prev_img = cv2.imread('assets/MOT17-mini/train/MOT17-2-FRCNN/img1/000001.jpg') | |
| warp_matrix = ecc.apply(prev_img, None) | |
| warp_matrix = ecc.apply(curr_img, None) | |
| start = time.process_time() | |
| for i in range(0, 100): | |
| warp_matrix = ecc.apply(prev_img, None) | |
| warp_matrix = ecc.apply(curr_img, None) | |
| end = time.process_time() | |
| print('Total time', end - start) | |
| print(warp_matrix) | |
| if ecc.prev_img_aligned is not None: | |
| curr_img = ecc.preprocess(curr_img) | |
| prev_img = ecc.preprocess(prev_img) | |
| weighted_img = cv2.addWeighted(curr_img, 0.5, ecc.prev_img_aligned, 0.5, 0) | |
| cv2.imshow('prev_img_aligned', weighted_img) | |
| cv2.waitKey(0) | |
| cv2.imwrite(str(BOXMOT / 'motion/cmc/ecc_aligned.jpg'), weighted_img) | |
| if __name__ == "__main__": | |
| main() | |