319 lines
10 KiB
Python
319 lines
10 KiB
Python
import cv2
|
|
import numpy as np
|
|
import imutils as im
|
|
|
|
from util import 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
|
|
|
|
CONSOLE_DEBUG = False
|
|
IMAGE_DEBUG = False
|
|
|
|
|
|
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 _debug(self, _image, matcher_res, matcher_start_x, matcher_start_y, template_scaled, crop_scaled):
|
|
template_img = self.ref['image']
|
|
offset = self.ref['offset']
|
|
|
|
(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_start : {(matcher_start_x, matcher_start_y)}")
|
|
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_img.shape}")
|
|
self._print(f"template_scaled : {template_scaled.shape}")
|
|
self._print(f"offset : {offset}")
|
|
|
|
if IMAGE_DEBUG:
|
|
matcher_bbox_local = (matcher_start_x, matcher_start_y, template_img.shape[1], template_img.shape[0])
|
|
matcher_crop = image_crop(_image, matcher_bbox_local)
|
|
cv2.imshow(f"{self.name}: Matcher res", matcher_res)
|
|
cv2.imshow(f"{self.name}: Matcher view", matcher_crop)
|
|
|
|
def _print(self, s):
|
|
if CONSOLE_DEBUG:
|
|
print(f"{self.name}: {s}")
|
|
|
|
def process(self, _image: np.array):
|
|
_image_anno = _image.copy()
|
|
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']
|
|
|
|
# Resize template
|
|
template_scaled = im.resize(template, IMG_SCALE_UP * template.shape[0], IMG_SCALE_UP * template.shape[1])
|
|
# Match
|
|
matcher_res = cv2.matchTemplate(crop_scaled, template_scaled, cv2.TM_CCOEFF_NORMED)
|
|
(min_val, max_val, min_loc, max_loc) = cv2.minMaxLoc(matcher_res)
|
|
|
|
# Get matcher coord
|
|
matcher_start_x = max_loc[0]/IMG_SCALE_UP
|
|
matcher_start_y = max_loc[1]/IMG_SCALE_UP
|
|
|
|
# Store offset at first frame
|
|
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)
|
|
|
|
# Compensate offset
|
|
matcher_start_x -= self.ref['offset'][0]
|
|
matcher_start_y -= self.ref['offset'][1]
|
|
|
|
matcher_bbox_local = (matcher_start_x, matcher_start_y, template.shape[1], template.shape[0])
|
|
|
|
self._debug(_image_anno, matcher_res, matcher_start_x, matcher_start_y, template_scaled, crop_scaled)
|
|
|
|
return matcher_bbox_local, matcher_res
|
|
|
|
def path_add(self, point):
|
|
if len(self.path) > 200:
|
|
return
|
|
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, color=(0, 255, 0), name: str = 'CornerTracker'):
|
|
self.color = color
|
|
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_ref = None
|
|
self.corner_matcher_list = []
|
|
|
|
def _print(self, s):
|
|
if CONSOLE_DEBUG:
|
|
print(f"{self.name}: {s}")
|
|
|
|
def _debug(self, _image_anno, tracking_anno, matcher_bbox_local):
|
|
matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb))
|
|
|
|
template_top_left, template_bottom_right = to_rect(self.matching_tpl_bb)
|
|
cv2.rectangle(tracking_anno, template_top_left, template_bottom_right, COLOR_TEMPLATE, 1)
|
|
|
|
matcher_top_left_local, matcher_bottom_right_local = to_rect(bbox_round(matcher_bbox_local))
|
|
cv2.rectangle(tracking_anno, matcher_top_left_local, matcher_bottom_right_local, COLOR_MATCHER, 1)
|
|
|
|
matcher_rect = to_rect(matcher_bbox)
|
|
cv2.rectangle(_image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)
|
|
|
|
tracker_rect = to_rect(self.tracking_bb)
|
|
cv2.rectangle(_image_anno, tracker_rect[0], tracker_rect[1], self.color, 1)
|
|
|
|
def init_reference_frame(self, _image: np.array):
|
|
print(f"Select tracking object")
|
|
bb = cv2.selectROI("Tracker Reference", _image, False)
|
|
cv2.destroyWindow("Tracker Reference")
|
|
if bbox_center(bb) == (0, 0):
|
|
return False
|
|
|
|
self.tracking_ref_bb = bb
|
|
self.tracking_ref_img = image_crop(_image, self.tracking_ref_bb)
|
|
self.tracking_ref_gray_img = cv2.cvtColor(self.tracking_ref_img, cv2.COLOR_BGR2GRAY)
|
|
self.corner_matcher_list = []
|
|
|
|
corner_list = []
|
|
count = 1
|
|
while True:
|
|
print(f"Add Corner {count}")
|
|
bb = cv2.selectROI("Matcher Reference", self.tracking_ref_img, fromCenter=True, showCrosshair=True)
|
|
if bbox_center(bb) == (0, 0):
|
|
break
|
|
|
|
self.matching_tpl_bb = bb
|
|
self.matching_tpl_img = image_crop(self.tracking_ref_gray_img.copy(), self.matching_tpl_bb)
|
|
self.corner_matcher_list.append(Corner(self.matching_tpl_img, self.matching_tpl_bb, name=f"Corner-{count}"))
|
|
corner_list.append(bbox_center(bbox_add_position(self.matching_tpl_bb, self.tracking_ref_bb)))
|
|
|
|
print(f"Corner {count} added")
|
|
print(f"Press any key to add another corner or ESC to continue")
|
|
|
|
count += 1
|
|
|
|
cv2.destroyWindow("Matcher Reference")
|
|
print(f"Added {count} corners")
|
|
|
|
if len(corner_list) == 0:
|
|
return False
|
|
|
|
# Initialize tracker with first frame and bounding box
|
|
self.tracker = cv2.TrackerKCF().create()
|
|
self.tracker.init(_image, self.tracking_ref_bb)
|
|
|
|
# Refine initial corners and store them as reference
|
|
self.corner_ref = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners=corner_list)
|
|
|
|
return True
|
|
|
|
def process(self, _image: np.array, _image_anno: 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, _image_anno)
|
|
|
|
def _match(self, _image: np.array, _image_anno: np.array):
|
|
tracking_anno = self.tracking_img.copy()
|
|
corners_raw = []
|
|
for _ct in self.corner_matcher_list:
|
|
matcher_bbox_local, matcher = _ct.process(self.tracking_img)
|
|
matcher_bbox = bbox_add_position(matcher_bbox_local, self.tracking_bb)
|
|
|
|
# Draw path
|
|
corner = bbox_center(matcher_bbox)
|
|
corners_raw.append(corner)
|
|
|
|
self._debug(_image_anno, tracking_anno, matcher_bbox_local)
|
|
|
|
# refine corners
|
|
corners_refined = self._corner_refine(cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY), corners_raw)
|
|
|
|
# store refined corners
|
|
_i = 0
|
|
for corner in corners_refined:
|
|
_ct = self.corner_matcher_list[_i]
|
|
_ct.path_add(corner)
|
|
_i += 1
|
|
|
|
# draw path
|
|
for _ct in self.corner_matcher_list:
|
|
for p in _ct.path:
|
|
cv2.line(_image_anno, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 1)
|
|
|
|
distances = []
|
|
for _i in range(0, corners_refined.shape[0]):
|
|
distances.append((corners_refined[_i] - self.corner_ref[_i]))
|
|
|
|
_mean_distance = np.mean(distances, axis=0)
|
|
|
|
# show refined corners
|
|
for _i in range(0, corners_refined.shape[0]):
|
|
self._print(f" -- Corner Reference [{i}] {self.corner_ref[_i]}")
|
|
self._print(f" -- Corner coarse [{i}] {corners_raw[_i]}")
|
|
self._print(f" -- Corner fine [{i}] {corners_refined[_i]}")
|
|
self._print(f" -- Corner distance [{i}] {corners_refined[_i] - self.corner_ref[_i]}")
|
|
cv2.circle(_image_anno, (int(corners_refined[_i, 0]), int(corners_refined[_i, 1])), 4, (0, 255, 0))
|
|
|
|
if IMAGE_DEBUG:
|
|
cv2.imshow(f"{self.name}: Tracker: ", self.tracking_img)
|
|
|
|
return _mean_distance
|
|
|
|
@staticmethod
|
|
def _corner_refine(src_gray, corners: np.array):
|
|
# convert from (n, 2) to (n, 1, 2)
|
|
corners_original = []
|
|
for corner in corners:
|
|
corners_original.append([corner])
|
|
|
|
corners_original = np.array(corners_original, dtype=np.float32)
|
|
|
|
# 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
|
|
corner_result = []
|
|
for _i in range(corners_original.shape[0]):
|
|
corner = (corners_refined[_i, 0, 0], corners_refined[_i, 0, 1])
|
|
corner_result.append(corner)
|
|
|
|
return np.array(corner_result)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
|
|
video = cv2.VideoCapture('./data/spindle_multi_black/%04d.png')
|
|
tracker_count = 0
|
|
tracker_list = []
|
|
if video.isOpened():
|
|
# Read first frame.
|
|
ok, image = video.read()
|
|
if ok:
|
|
for tracker_count in range(0, 6):
|
|
select_window = image.copy()
|
|
print(f"Add tracker #{tracker_count}")
|
|
ct = CornerTracker(colors[tracker_count], name=f"Tracker-{tracker_count}")
|
|
ok = ct.init_reference_frame(select_window)
|
|
if ok:
|
|
tracker_list.append(ct)
|
|
print(f"Number of active tracker: {len(tracker_list)}")
|
|
else:
|
|
break
|
|
|
|
else:
|
|
print('Cannot read video file')
|
|
|
|
key_wait = -1
|
|
while True:
|
|
# Read a new frame
|
|
ok, image = video.read()
|
|
if ok:
|
|
image_anno = image.copy()
|
|
i = 0
|
|
for ct in tracker_list:
|
|
mean_distance = ct.process(image, image_anno)
|
|
if mean_distance is not None:
|
|
cv2.putText(image_anno, f"Distance [{i}] : ({mean_distance[0]:+05.2f}, {mean_distance[1]:+05.2f})",
|
|
(20, 25 + 25 * i), cv2.FONT_HERSHEY_SIMPLEX, 0.5, ct.color, 1)
|
|
i += 1
|
|
cv2.imshow(f"Image Anno", image_anno)
|
|
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 = 1
|
|
else:
|
|
key_wait = -1
|
|
|
|
else:
|
|
print('Cannot open video file')
|