ocv_template_matching/ocv_corner_tracker.py

import cv2
import numpy as np
import imutils as im

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

IMG_SCALE_UP = 1
IMG_ROTATE = 0

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


class Corner:
	def __init__(self, reference: np.array, bbox: np.array, name: str = 'Corner'):
		self.ref = {'image': reference, 'bbox': bbox, 'offset': (0, 0)}
		self.name = name
		self._path = []
		self.line_from = None
		self.first_frame = True

	def _print(self, s):
		print(f"{self.name}: {s}")

	def process(self, image: np.array):
		crop_scaled = im.resize(image, IMG_SCALE_UP * image.shape[1], IMG_SCALE_UP * image.shape[0])
		template = self.ref['image']
		bb_template = self.ref['bbox']

		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)
		self._print(f"(min_val, max_val, min_loc, max_loc): {(min_val, max_val, min_loc, max_loc)}")
		self._print(f"matcher_res     : {matcher_res.shape}")
		self._print(f"crop            : {image.shape}")
		self._print(f"crop_scaled     : {crop_scaled.shape}")
		self._print(f"template        : {template.shape}")
		self._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 self.first_frame:
			self.first_frame = False
			offset_x = matcher_start_x - bb_template[0]
			offset_y = matcher_start_y - bb_template[1]
			self.ref['offset'] = (offset_x, offset_y)
			self._print(f"offset : {(offset_x, offset_y)}")

		matcher_start_x -= self.ref['offset'][0]
		matcher_start_y -= self.ref['offset'][1]

		self._print(f"matcher_start : {(matcher_start_x, matcher_start_y)}")
		matcher_bbox_local = (matcher_start_x, matcher_start_y, template.shape[1], template.shape[0])

		cv2.imshow(f"{self.name}: Matcher res", matcher_res)

		return matcher_bbox_local, matcher_res

	def path_add(self, point):
		if self.line_from is not None:
			self._path.append({'from': self.line_from, 'to': point})
		self.line_from = point

	@property
	def path(self):
		return self._path


class CornerTracker:
	def __init__(self, name: str = 'CornerTracker'):
		self.name = name
		self.tracking_ref_bb = None
		self.tracking_ref_img = None
		self.tracking_ref_gray_img = None
		self.matching_tpl_bb = None
		self.matching_tpl_img = None
		self.tracking_bb = None
		self.tracking_img = None
		self.tracker = None
		self.corner_list = []

	def _print(self, s):
		print(f"{self.name}: {s}")

	def _debug(self, image_anno, matcher_bbox_local):
		matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb))
		matcher_crop = image_crop(image, matcher_bbox)
		cv2.imshow("{self.name}: Matcher view", matcher_crop)

		matcher_rect = to_rect(matcher_bbox)
		tracker_rect = to_rect(self.tracking_bb)
		template_top_left, template_bottom_right = to_rect(self.matching_tpl_bb)
		matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local))

		cv2.rectangle(self.tracking_img, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1)
		cv2.rectangle(self.tracking_img, 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)

	def init_reference_frame(self, image: np.array):
		print(f"Select tracking object")
		self.tracking_ref_bb = cv2.selectROI("Image", image, False)
		self.tracking_ref_img = image_crop(image.copy(), self.tracking_ref_bb)
		self.tracking_ref_gray_img = cv2.cvtColor(self.tracking_ref_img, cv2.COLOR_BGR2GRAY)

		self.corner_list = []
		count = 1
		while True:
			self._print(f"Add Corner {count}")
			self.matching_tpl_bb = cv2.selectROI("Image", self.tracking_ref_img, fromCenter=True, showCrosshair=True)
			self.matching_tpl_img = image_crop(self.tracking_ref_gray_img.copy(), self.matching_tpl_bb)
			self.corner_list.append(Corner(self.matching_tpl_img, self.matching_tpl_bb))
			self._print(f"Corner {count} added")
			self._print(f"Press any key to add another corner or ESC to continue")

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

		self._print(f"Added {count} corners")

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

	def process(self, image: np.array):
		if self.tracker is None:
			raise Exception(f"{self.name}: Call init_reference_frame() first")

		# Update tracker
		_ok, tracking_bb = self.tracker.update(image)
		if not _ok:
			return None

		tracking_img = image_crop(image.copy(), tracking_bb)
		tracking_img = cv2.cvtColor(tracking_img, cv2.COLOR_BGR2GRAY)
		self.tracking_img = cv2.GaussianBlur(tracking_img, (9, 9), 0)
		self.tracking_bb = tracking_bb

		return self._match(image)

	def _match(self, image: np.array):
		image_anno = image.copy()
		corners = []
		for ct in self.corner_list:
			matcher_bbox_local, matcher = ct.process(self.tracking_img)
			matcher_bbox = bbox_add_position(matcher_bbox_local, self.tracking_bb)

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

			self._debug(image_anno, matcher_bbox_local)

		# refine corners
		corners_refined = self._corner_refine(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), np.array(corners, dtype=np.float32))

		# store refined corners
		for i in range(corners_refined.shape[0]):
			line_to = (corners_refined[i, 0, 0], corners_refined[i, 0, 1])
			ct = self.corner_list[i]
			ct.path_add(line_to)

		# draw path
		for ct in self.corner_list:
			for p in ct.path:
				cv2.line(image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 1)

		# show refined corners
		for i in range(corners_refined.shape[0]):
			self._print(f" -- Refined Corner [{i}] ({corners_refined[i, 0, 0]}, {corners_refined[i, 0, 1]})")
			cv2.circle(image_anno, (int(corners_refined[i, 0, 0]), int(corners_refined[i, 0, 1])), 4, (0, 255, 0))

		cv2.imshow(f"{self.name}: Tracker: ", self.tracking_img)
		cv2.imshow(f"{self.name}: Image Anno", image_anno)

	def _corner_refine(self, 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]):
			self._print(f" -- Original Corner [{i}] ({corners_original[i, 0, 0]}, {corners_original[i, 0, 1]})")
			self._print(f" -- Refined Corner  [{i}] ({corners_refined[i, 0, 0]}, {corners_refined[i, 0, 1]})")

		return corners_refined


if __name__ == '__main__':
	video = cv2.VideoCapture('./data/spindle_multi_black/%04d.png')
	if video.isOpened():
		ct = CornerTracker()
		# Read first frame.
		ok, image = video.read()
		if ok:
			ct.init_reference_frame(image)
		else:
			print('Cannot read video file')

		key_wait = -1
		while True:
			# Read a new frame
			ok, image = video.read()
			if ok:
				ct.process(image)
			else:
				video.set(cv2.CAP_PROP_POS_FRAMES, 0)
				continue

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

	else:
		print('Cannot open video file')