python
复制代码
import copy
import time
import cv2
import numpy as np
def draw_kpts(image0, image1, mkpts0, mkpts1, margin=10):
H0, W0 = image0.shape
H1, W1 = image1.shape
H, W = max(H0, H1), W0 + W1 + margin
out = 255 * np.ones((H, W), np.uint8)
out[:H0, :W0] = image0
out[:H1, W0+margin:] = image1
out = np.stack([out]*3, -1)
mkpts0, mkpts1 = np.round(mkpts0).astype(int), np.round(mkpts1).astype(int)
# print(f"mkpts0.shape : {mkpts0.shape}")
c = (0, 255, 0)
for (new, old) in zip(mkpts0, mkpts1):
x0, y0 = new.ravel()
x1, y1 = old.ravel()
# print(f"x0 : {x0}")
# cv2.line(out, (x0, y0), (x1 + margin + W0, y1),
# color=c, thickness=1, lineType=cv2.LINE_AA)
# display line end-points as circles
cv2.circle(out, (x0, y0), 2, c, -1, lineType=cv2.LINE_AA)
cv2.circle(out, (x1 + margin + W0, y1), 2, c, -1,
lineType=cv2.LINE_AA)
return out
if __name__ == "__main__":
img0Path = "/training/datasets/orchard/orchard_imgs_/000130.jpg"
img1Path = "/training/datasets/orchard/orchard_imgs_/000132.jpg"
img0 = cv2.imread(img0Path, 0)
img1 = cv2.imread(img1Path, 0)
h, w = img0.shape
mask = np.zeros_like(img0)
mask[int(0.02 * h): int(0.98 * h), int(0.02 * w): int(0.98 * w)] = 255
keypoints = cv2.goodFeaturesToTrack(
img0,
mask=mask,
maxCorners=2048,
qualityLevel=0.01,
minDistance=1,
blockSize=3,
useHarrisDetector=False,
k=0.04
)
print(f"keypoints , size : {keypoints.shape}")
next_keypoints, status, err = cv2.calcOpticalFlowPyrLK(
img0, img1, keypoints, None
)
H, _ = cv2.estimateAffinePartial2D(
keypoints, next_keypoints, cv2.RANSAC
)
print(f"H : {H}")
out = draw_kpts(img0, img1, keypoints , next_keypoints)
cv2.imwrite("keypoints.jpg", out)
