ocv_template_matching/ocv_template_matching.py

import cv2
import sys
import imutils as im
import numpy as np

from util import bbox_extend, bbox_round, image_crop, to_rect, bbox_add_position, bbox_center
from ocv_corner_tracker import CornerTracker


def refine_corner(src_gray, corners_original: np.array):
	# Set the needed parameters to find the refined corners
	win_size = (5, 5)
	zero_zone = (-1, -1)
	criteria = (cv2.TERM_CRITERIA_EPS + cv2.TermCriteria_COUNT, 40, 0.001)

	# Calculate the refined corner locations
	corners_refined = cv2.cornerSubPix(src_gray, corners_original.copy(), win_size, zero_zone, criteria)

	# Write them down
	for i in range(corners_original.shape[0]):
		print(" -- Original Corner [", i, "] (", corners_original[i, 0, 0], ",", corners_original[i, 0, 1], ")")
		print(" -- Refined Corner  [", i, "] (", corners_refined[i, 0, 0], ",", corners_refined[i, 0, 1], ")")


COLOR_TRACKER = (0, 255, 0)
COLOR_TRACKER_EXT = (0, 255, 255)
COLOR_MATCHER = (0, 0, 255)
COLOR_TEMPLATE = (255, 0, 0)
IMG_SCALE_UP = 1
TEMPLATE_MATCH_OVERLAP = 0

#video = cv2.VideoCapture("./data/IMG_1730.MP4")
video = cv2.VideoCapture('./data/spindle_multi_black/%04d.png')

# Exit if video not open
# d.
if not video.isOpened():
	print("Could not open video")
	sys.exit()

# Read first frame.
ok, image = video.read()
if not ok:
	print('Cannot read video file')
	sys.exit()

bb_crop = bbox_extend(cv2.selectROI("Image", image, False), TEMPLATE_MATCH_OVERLAP)
crop = image_crop(image.copy(), bb_crop)

ct_list = []
count = 1
key_wait = -1

while True:
	print(f"Add Corner {count}")
	bb_template = cv2.selectROI("Image", crop, fromCenter=True, showCrosshair=True)
	template = cv2.cvtColor(image_crop(crop.copy(), bb_template), cv2.COLOR_BGR2GRAY)
	ct_list.append(CornerTracker(template, bb_template))
	print(f"Corner {count} added")
	print(f"Press any key to add another corner or ESC to continue")

	k = cv2.waitKey(-1) & 0xff
	if k == 27:
		break
	count += 1

print(f"Added {count} corners")

# Initialize tracker with first frame and bounding box
tracker = cv2.TrackerKCF().create()
tracker.init(image, bb_crop)

path = []
line_from = None

is_startup = True
offset_x = 0
offset_y = 1
mode = True
while True:
	# Read a new frame
	ok, image = video.read()
	if not ok:
		break

	# Start timer
	timer = cv2.getTickCount()

	# Update tracker
	ok, bb_crop = tracker.update(image)
	crop = image_crop(image.copy(), bb_crop)
	crop = cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)
	crop = cv2.GaussianBlur(crop, (15, 15), 0)
	image_anno = image.copy()

	if ok:
		if mode:
			corners = []
			for ct in ct_list:
				matcher_bbox_local, matcher = ct.process(crop)
				matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local))

				matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, bb_crop))
				matcher_rect = to_rect(matcher_bbox)

				tracker_rect = to_rect(bb_crop)
				template_top_left, template_bottom_right = to_rect(bb_template)

				cv2.rectangle(crop, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1)
				cv2.rectangle(crop, matcher_top_left_local, matcher_bottom_right_local, COLOR_MATCHER, 1)

				cv2.rectangle(image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)
				cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], COLOR_TRACKER, 1)
				matcher_crop = image_crop(image, matcher_bbox)

				# Draw path
				line_to = bbox_center(matcher_bbox)
				corners.append([line_to])
				ct.path_add(line_to)

				for p in ct.path:
					cv2.line(image_anno, p['from'], p['to'], (0, 0, 255), 2)

			refine_corner(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), np.array(corners, dtype=np.float32))

		else:
			# Perform template match operations
			crop_scaled = im.resize(crop.copy(), IMG_SCALE_UP * crop.shape[1], IMG_SCALE_UP * crop.shape[0])
			template_scaled = im.resize(template, IMG_SCALE_UP * template.shape[0], IMG_SCALE_UP * template.shape[1])
			matcher_res = cv2.matchTemplate(crop_scaled, template_scaled, cv2.TM_CCOEFF_NORMED)
			(min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(matcher_res)
			print(f"(min_val, max_val, min_loc, max_loc): {(min_val, max_val, min_loc, max_loc)}")
			print(f"matcher_res     : {matcher_res.shape}")
			print(f"crop     		: {crop.shape}")
			print(f"crop_scaled     : {crop_scaled.shape}")
			print(f"template        : {template.shape}")
			print(f"template_scaled : {template_scaled.shape}")

			matcher_start_x = max_loc[0]/IMG_SCALE_UP
			matcher_start_y = max_loc[1]/IMG_SCALE_UP
			if is_startup:
				is_startup = False
				offset_x = matcher_start_x - bb_template[0]
				offset_y = matcher_start_y - bb_template[1]
				print(f"offset : {(offset_x, offset_y)}")

			matcher_start_x -= offset_x
			matcher_start_y -= offset_y
			print(f"matcher_start : {(matcher_start_x, matcher_start_y)}")
			matcher_bbox_local = (round(matcher_start_x), round(matcher_start_y), template.shape[1], template.shape[0])
			matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local))

			matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, bb_crop))
			matcher_rect = to_rect(matcher_bbox)

			tracker_rect = to_rect(bb_crop)
			template_top_left, template_bottom_right = to_rect(bb_template)

			cv2.rectangle(crop, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1)
			cv2.rectangle(crop, matcher_top_left_local, matcher_bottom_right_local, COLOR_MATCHER, 1)

			cv2.rectangle(image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)
			cv2.rectangle(image_anno, tracker_rect[0], tracker_rect[1], COLOR_TRACKER, 1)
			matcher_crop = image_crop(image, matcher_bbox)

			# Draw path
			line_to = bbox_center(matcher_bbox)
			if line_from is not None:
				path.append({'from': line_from, 'to': line_to})
			line_from = line_to

			for p in path:
				cv2.line(image_anno, p['from'], p['to'], (0, 0, 255), 2)

			cv2.imshow("Matcher res", matcher_res)
			cv2.imshow("Matcher view", matcher_crop)

		cv2.imshow("Crop", crop)
		cv2.imshow("Image", image_anno)

	# Exit if ESC pressed
	k = cv2.waitKey(key_wait) & 0xff
	if k == 27:
		break
	if k == ord('g'):
		key_wait = 30
	if k == ord('s'):
		key_wait = -1