117 lines
4.4 KiB
Python
117 lines
4.4 KiB
Python
import cv2 as cv
|
|
import dlib
|
|
import numpy as np
|
|
import os
|
|
|
|
def match_feature_find_object(query_img, train_img, min_matches, results_file_name):
|
|
# Create an ORB object
|
|
orb = cv.ORB_create(nfeatures=100000)
|
|
|
|
features1, des1 = orb.detectAndCompute(query_img, None)
|
|
features2, des2 = orb.detectAndCompute(train_img, None)
|
|
|
|
# Create Brute-Force matcher object
|
|
bf = cv.BFMatcher(cv.NORM_HAMMING)
|
|
matches = bf.knnMatch(des1, des2, k=2)
|
|
|
|
# Nearest neighbour ratio test to find good matches
|
|
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) >= min_matches:
|
|
print(f"good: {len(good)}, {results_file_name}")
|
|
# Draw a polygon around the recognized object
|
|
src_pts = np.float32([features1[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)
|
|
|
|
# Get the transformation matrix
|
|
M, _ = cv.findHomography(src_pts, dst_pts, cv.RANSAC, 5.0)
|
|
|
|
# Find the perspective transformation to get the corresponding points
|
|
h, w = query_img.shape[:2]
|
|
pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)
|
|
dst = cv.perspectiveTransform(pts, M)
|
|
|
|
train_img = cv.polylines(train_img, [np.int32(dst)], True, (0, 255, 0), 2, cv.LINE_AA)
|
|
|
|
# Return the bounding box of the detected object
|
|
x, y, w, h = cv.boundingRect(np.int32(dst))
|
|
return (x, y, w, h), train_img # Return bounding box and modified image
|
|
|
|
else:
|
|
print("Not enough good matches are found - {}/{}".format(len(good), min_matches))
|
|
return None, train_img # No detection, return None
|
|
|
|
|
|
def main():
|
|
input_folder = "extracted_frames/"
|
|
output_folder = "results_extracted_frames/"
|
|
new_object_size = (300, 300)
|
|
|
|
os.makedirs(input_folder, exist_ok=True)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
orb = cv.ORB_create()
|
|
query_img = cv.imread("img/drone.png")
|
|
query_img = cv.resize(query_img, new_object_size)
|
|
features = orb.detect(query_img, None)
|
|
f_img = cv.drawKeypoints(query_img, features, None, color=(0, 255, 0), flags=0)
|
|
# cv.imwrite(os.path.join(output_folder, "drone.png"), f_img) # Save the image with keypoints
|
|
|
|
# Initialize dlib correlation tracker
|
|
tracker = dlib.correlation_tracker()
|
|
is_tracking = False # Flag to indicate if we are tracking
|
|
|
|
# Sort the filenames to process them in sequence
|
|
filenames = sorted([f for f in os.listdir(input_folder) if f.endswith("png")])
|
|
|
|
# Previous bounding box (if the object is detected)
|
|
prev_bbox = None
|
|
|
|
for filename in filenames:
|
|
img = cv.imread(os.path.join(input_folder, filename))
|
|
|
|
if not is_tracking: # Try to detect using feature matching
|
|
bbox, modified_img = match_feature_find_object(query_img, img, 100, os.path.join(output_folder, filename))
|
|
|
|
if bbox is not None:
|
|
# Object detected, initialize the dlib tracker
|
|
x, y, w, h = bbox
|
|
tracker.start_track(img, dlib.rectangle(x, y, x + w, y + h))
|
|
is_tracking = True
|
|
prev_bbox = bbox
|
|
else:
|
|
# No detection, use the image as is
|
|
modified_img = img
|
|
|
|
else: # If we are tracking
|
|
# Update the tracker and get the new position
|
|
tracker.update(img)
|
|
pos = tracker.get_position()
|
|
x1, y1, x2, y2 = int(pos.left()), int(pos.top()), int(pos.right()), int(pos.bottom())
|
|
print(x1, y1, x2, y2)
|
|
|
|
# Draw the tracking bounding box
|
|
cv.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# If the bounding box is valid, keep tracking
|
|
if (x2 - x1) > 0 and (y2 - y1) > 0:
|
|
prev_bbox = (x1, y1, x2 - x1, y2 - y1)
|
|
modified_img = img
|
|
else:
|
|
# If tracking fails, stop tracking and try detection again in the next frame
|
|
is_tracking = False
|
|
modified_img = img
|
|
|
|
# Save the result as a PNG image
|
|
# cv.imwrite(os.path.join(output_folder, filename), modified_img)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|