File size: 1,530 Bytes
3a66575 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | import cv2
import numpy as np
class GlobalMotionEstimator:
def __init__(self):
self.prev_gray = None
self.prev_pts = None
def update(self, current_frame):
# Expects a BGR or RGB frame (function works identically long as it's 3-channel)
curr_gray = cv2.cvtColor(current_frame, cv2.COLOR_BGR2GRAY)
transform = np.eye(2, 3, dtype=np.float32)
if self.prev_gray is None:
self.prev_gray = curr_gray
self.prev_pts = cv2.goodFeaturesToTrack(curr_gray, maxCorners=200, qualityLevel=0.01, minDistance=30, blockSize=3)
return transform
if self.prev_pts is not None and len(self.prev_pts) > 0:
curr_pts, status, err = cv2.calcOpticalFlowPyrLK(self.prev_gray, curr_gray, self.prev_pts, None, winSize=(21, 21), maxLevel=3)
good_prev = self.prev_pts[status == 1]
good_curr = curr_pts[status == 1]
if len(good_prev) >= 4 and len(good_curr) >= 4:
# Estimate affine transform (translation + rotation + scale)
transform, inliers = cv2.estimateAffinePartial2D(good_prev, good_curr)
if transform is None: # Safety fallback
transform = np.eye(2, 3, dtype=np.float32)
self.prev_gray = curr_gray
self.prev_pts = cv2.goodFeaturesToTrack(curr_gray, maxCorners=200, qualityLevel=0.01, minDistance=30, blockSize=3)
return transform
|