50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
import os
|
|
import cv2 as cv
|
|
from detector import Detector
|
|
import matplotlib.pyplot as plt
|
|
|
|
SHOW_FAILED = False
|
|
|
|
def show_image(img, title, colorspace, results_file_name, show=False):
|
|
if not show:
|
|
return
|
|
dpi = 96
|
|
figsize = (img.shape[1] / dpi, img.shape[0] / dpi)
|
|
fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
|
|
if colorspace == "RGB":
|
|
plt.imshow(cv.cvtColor(img, cv.COLOR_BGR2RGB), interpolation="spline16")
|
|
if colorspace == "gray":
|
|
plt.imshow(img, cmap="gray")
|
|
plt.title(title, fontsize=12)
|
|
ax.axis("off")
|
|
plt.savefig(results_file_name)
|
|
plt.close()
|
|
|
|
def main():
|
|
input_folder = "data/video-frames/"
|
|
success_output_folder = "data/success-video-frames/"
|
|
failures_output_folder = "data/failures-video-frames/"
|
|
|
|
os.makedirs(input_folder, exist_ok=True)
|
|
os.makedirs(success_output_folder, exist_ok=True)
|
|
os.makedirs(failures_output_folder, exist_ok=True)
|
|
|
|
target_image = cv.imread("cropped_image.png")
|
|
|
|
detector = Detector(target_image, min_matches=4, angle_tolerance=20, length_ratio_tolerance=2.5)
|
|
|
|
filenames = sorted([f for f in os.listdir(input_folder) if f.endswith("png")])
|
|
|
|
for filename in filenames:
|
|
img = cv.imread(os.path.join(input_folder, filename))
|
|
bbox, modified_img = detector.match_feature_find_object(img)
|
|
if bbox is not None:
|
|
show_image(modified_img, "Success feature matching", "RGB", os.path.join(success_output_folder, filename), show=True)
|
|
else:
|
|
show_image(modified_img, "Failed feature matching", "RGB", os.path.join(failures_output_folder, filename), show=SHOW_FAILED)
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|