73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
import os
|
|
import cv2 as cv
|
|
import numpy as np
|
|
from detector import Detector
|
|
from tracker import Tracker
|
|
import matplotlib.pyplot as plt
|
|
|
|
|
|
def show_image(img, title, colorspace, results_file_name, show=False):
|
|
"""Displays or saves the processed image."""
|
|
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 draw_bounding_box(img, bbox, color=(0, 255, 0)):
|
|
"""Draws a bounding box (as a polygon) on the image."""
|
|
x, y, w, h = bbox
|
|
box_points = np.array([
|
|
[x, y],
|
|
[x + w, y],
|
|
[x + w, y + h],
|
|
[x, y + h]
|
|
], dtype=np.int32)
|
|
cv.polylines(img, [box_points], isClosed=True, color=color, thickness=2, lineType=cv.LINE_AA)
|
|
|
|
|
|
def main():
|
|
input_folder = "data/video-frames/"
|
|
output_folder = "data/detect_and_track/results/"
|
|
|
|
os.makedirs(input_folder, exist_ok=True)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
|
|
query_img = cv.imread("cropped_image.png")
|
|
detector = Detector(query_img, min_matches=4, angle_tolerance=20, length_ratio_tolerance=2.5)
|
|
tracker = Tracker()
|
|
|
|
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, detected_img = detector.match_feature_find_object(img)
|
|
if bbox is not None:
|
|
tracker.start_tracking(img, bbox)
|
|
draw_bounding_box(img, bbox, color=(0, 255, 0))
|
|
modified_img = detected_img
|
|
else:
|
|
if tracker.is_tracking:
|
|
bbox = tracker.update(img)
|
|
if bbox:
|
|
draw_bounding_box(img, bbox, color=(255, 0, 0))
|
|
else:
|
|
tracker.is_tracking = False
|
|
modified_img = img
|
|
|
|
show_image(modified_img, "Feature matching and object recognition", "RGB", os.path.join(output_folder, filename), show=(bbox is not None))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|