Tracker: clean up GMC (#8003)

pull/8018/head
Laughing 1 year ago committed by GitHub
parent afc5dcb3f9
commit c26e497eec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 16
      ultralytics/trackers/utils/gmc.py

@ -97,13 +97,13 @@ class GMC:
if self.method in ["orb", "sift"]:
return self.applyFeatures(raw_frame, detections)
elif self.method == "ecc":
return self.applyEcc(raw_frame, detections)
return self.applyEcc(raw_frame)
elif self.method == "sparseOptFlow":
return self.applySparseOptFlow(raw_frame, detections)
return self.applySparseOptFlow(raw_frame)
else:
return np.eye(2, 3)
def applyEcc(self, raw_frame: np.array, detections: list = None) -> np.array:
def applyEcc(self, raw_frame: np.array) -> np.array:
"""
Apply ECC algorithm to a raw frame.
@ -144,7 +144,7 @@ class GMC:
# 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:
(cc, H) = cv2.findTransformECC(self.prevFrame, frame, H, self.warp_mode, self.criteria, None, 1)
(_, 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}")
@ -275,7 +275,7 @@ class GMC:
# plt.show()
# Find rigid matrix
if (prevPoints.shape[0] > 4) and (prevPoints.shape[0] == prevPoints.shape[0]):
if prevPoints.shape[0] > 4:
H, inliers = cv2.estimateAffinePartial2D(prevPoints, currPoints, cv2.RANSAC)
# Handle downscale
@ -292,7 +292,7 @@ class GMC:
return H
def applySparseOptFlow(self, raw_frame: np.array, detections: list = None) -> np.array:
def applySparseOptFlow(self, raw_frame: np.array) -> np.array:
"""
Apply Sparse Optical Flow method to a raw frame.
@ -328,7 +328,7 @@ class GMC:
return H
# Find correspondences
matchedKeypoints, status, err = cv2.calcOpticalFlowPyrLK(self.prevFrame, frame, self.prevKeyPoints, None)
matchedKeypoints, status, _ = cv2.calcOpticalFlowPyrLK(self.prevFrame, frame, self.prevKeyPoints, None)
# Leave good correspondences only
prevPoints = []
@ -344,7 +344,7 @@ class GMC:
# Find rigid matrix
if (prevPoints.shape[0] > 4) and (prevPoints.shape[0] == prevPoints.shape[0]):
H, inliers = cv2.estimateAffinePartial2D(prevPoints, currPoints, cv2.RANSAC)
H, _ = cv2.estimateAffinePartial2D(prevPoints, currPoints, cv2.RANSAC)
if self.downscale > 1.0:
H[0, 2] *= self.downscale

Loading…
Cancel
Save