82 lines
3.2 KiB
Python
82 lines
3.2 KiB
Python
import cv2 as cv
|
|
import numpy as np
|
|
|
|
class Detector:
|
|
def __init__(self, target_image, min_matches=4, angle_tolerance=30, length_ratio_tolerance=2.5, nfeatures=30000):
|
|
self.orb = cv.ORB_create(nfeatures=nfeatures)
|
|
self.target_features, self.target_des = self.orb.detectAndCompute(target_image, None)
|
|
self.target_image_h, self.target_image_w = target_image.shape[:2]
|
|
self.min_matches = min_matches
|
|
self.angle_tolerance = angle_tolerance
|
|
self.length_ratio_tolerance = length_ratio_tolerance
|
|
|
|
def valid_rectangle(self, points):
|
|
def angle_between(p1, p2, p3):
|
|
v1 = p1 - p2
|
|
v2 = p3 - p2
|
|
angle = np.arccos(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
|
|
return np.degrees(angle)
|
|
|
|
if len(points) != 4:
|
|
return False
|
|
|
|
angles = [
|
|
angle_between(points[i], points[(i + 1) % 4], points[(i + 2) % 4])
|
|
for i in range(4)
|
|
]
|
|
if not all(abs(angle - 90) <= self.angle_tolerance for angle in angles):
|
|
print(f'Not a valid rectangle, angle criterion {[angle for angle in angles]}')
|
|
return False
|
|
|
|
side_lengths = [
|
|
np.linalg.norm(points[i] - points[(i + 1) % 4]) for i in range(4)
|
|
]
|
|
max_length = max(side_lengths)
|
|
min_length = min(side_lengths)
|
|
if max_length / min_length > self.length_ratio_tolerance:
|
|
print(f'Not a valid rectangle, length criterion')
|
|
return False
|
|
|
|
return True
|
|
|
|
def match_feature_find_object(self, train_img):
|
|
features2, des2 = self.orb.detectAndCompute(train_img, None)
|
|
|
|
bf = cv.BFMatcher(cv.NORM_HAMMING)
|
|
matches = bf.knnMatch(self.target_des, des2, k=2)
|
|
|
|
good = []
|
|
good_without_lists = []
|
|
matches = [match for match in matches if len(match) == 2]
|
|
for m, n in matches:
|
|
if m.distance < 0.8 * n.distance:
|
|
good.append([m])
|
|
good_without_lists.append(m)
|
|
|
|
if len(good) >= self.min_matches:
|
|
src_pts = np.float32([self.target_features[m.queryIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)
|
|
dst_pts = np.float32([features2[m.trainIdx].pt for m in good_without_lists]).reshape(-1, 1, 2)
|
|
|
|
M, _ = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)
|
|
|
|
h, w = self.target_image_h, self.target_image_w
|
|
pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
|
|
|
|
try:
|
|
dst = cv.perspectiveTransform(pts, M)
|
|
except cv.error as e:
|
|
print(f"Perspective transform failed: {e}")
|
|
return None, train_img
|
|
|
|
rectangle_points = np.squeeze(dst)
|
|
if not self.valid_rectangle(rectangle_points):
|
|
return None, train_img
|
|
|
|
train_img = cv.polylines(train_img, [np.int32(dst)], True, (0, 255, 0), 2, cv.LINE_AA)
|
|
x, y, w, h = cv.boundingRect(np.int32(dst))
|
|
return (x, y, w, h), train_img
|
|
|
|
else:
|
|
print(f"Not enough good matches are found - {len(good)}/{self.min_matches}")
|
|
return None, train_img
|