386 lines
12 KiB
Python
386 lines
12 KiB
Python
import cv2
|
|
import numpy as np
|
|
import argparse
|
|
import sys
|
|
|
|
from util import image_crop, to_rect, bbox_add_position, bbox_center, bbox_round
|
|
|
|
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)
|
|
COLOR_TRACK = (255, 255, 0)
|
|
COLOR_CIRCLE = (0, 255, 0)
|
|
|
|
CONSOLE_DEBUG = False
|
|
|
|
|
|
class Corner:
|
|
def __init__(self, _ref_img: np.array, _ref_bb: np.array, name: str = 'Corner'):
|
|
self.ref_img = _ref_img
|
|
self.ref_bb = _ref_bb
|
|
self.ref_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):
|
|
(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"template : {self.ref_img.shape}")
|
|
self._print(f"offset : {self.ref_offset}")
|
|
|
|
def _print(self, s):
|
|
if CONSOLE_DEBUG:
|
|
print(f"{self.name}: {s}")
|
|
|
|
def process(self, _image: np.array):
|
|
# Match
|
|
matcher_res = cv2.matchTemplate(_image, self.ref_img, 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 - self.ref_bb[0]
|
|
offset_y = matcher_start_y - self.ref_bb[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, self.ref_img.shape[1], self.ref_img.shape[0])
|
|
|
|
self._debug(_image, matcher_res, matcher_start_x, matcher_start_y)
|
|
|
|
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 CornerTrackerParams:
|
|
def __init__(self, _scale: float = 1.0, _pre_track: bool = False, _show_path: bool = False):
|
|
self.scale = _scale
|
|
self.pre_track = _pre_track
|
|
self.show_path = _show_path
|
|
|
|
|
|
class CornerTracker:
|
|
def __init__(self, _params: CornerTrackerParams, color=(0, 255, 0), name: str = 'CornerTracker'):
|
|
self.params = _params
|
|
self.color = color
|
|
self.name = name
|
|
self.matching_tpl_bb = None
|
|
self.tracking_bb = None
|
|
self.tracking_mask = 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, matcher_bbox_local):
|
|
matcher_bbox = bbox_round(bbox_add_position(matcher_bbox_local, self.tracking_bb))
|
|
|
|
matcher_rect = to_rect(matcher_bbox)
|
|
cv2.rectangle(_image_anno, matcher_rect[0], matcher_rect[1], COLOR_MATCHER, 1)
|
|
|
|
tracker_rect = to_rect(bbox_round(self.tracking_bb))
|
|
cv2.rectangle(_image_anno, tracker_rect[0], tracker_rect[1], self.color, 1)
|
|
|
|
@staticmethod
|
|
def mask_init(_image: np.array):
|
|
# Create mask
|
|
_mask = np.zeros((_image.shape[0], _image.shape[1]), dtype=np.uint8)
|
|
return _mask
|
|
|
|
@staticmethod
|
|
def mask_draw(_image: np.array, _mask: np.array):
|
|
print(f"Draw mask object")
|
|
bb = cv2.selectROI("Tracker Reference", _image, False)
|
|
cv2.destroyWindow("Tracker Reference")
|
|
if bbox_center(bb) == (0, 0):
|
|
bb = (0, 0, _image.shape[1], _image.shape[0])
|
|
|
|
rect_upper_left = (bb[0], bb[1])
|
|
rect_lower_right = (bb[0] + bb[2], bb[1] + bb[3])
|
|
cv2.rectangle(_mask, rect_upper_left, rect_lower_right, 255, -1)
|
|
|
|
return True
|
|
|
|
@staticmethod
|
|
def mask_apply(_image: np.array, _mask: np.array):
|
|
_image[_mask == 0] = 0
|
|
_image[_mask != 0] = _image[_mask != 0]
|
|
|
|
@staticmethod
|
|
def image_process(_image: np.array):
|
|
result = cv2.cvtColor(_image, cv2.COLOR_BGR2GRAY)
|
|
|
|
return result
|
|
|
|
def init_reference_frame(self, _image: np.array):
|
|
# Draw mask
|
|
_mask = CornerTracker.mask_init(_image)
|
|
CornerTracker.mask_draw(_image, _mask)
|
|
|
|
# Select search area
|
|
print(f"Select tracking object")
|
|
bb = cv2.selectROI("Tracker Reference", _image, False)
|
|
cv2.destroyWindow("Tracker Reference")
|
|
if bbox_center(bb) == (0, 0):
|
|
return False
|
|
|
|
tracking_img = image_crop(_image, bb)
|
|
tracking_mask = image_crop(_mask, bb)
|
|
CornerTracker.mask_apply(tracking_img, tracking_mask)
|
|
image_processed_local = CornerTracker.image_process(tracking_img)
|
|
|
|
self.corner_matcher_list = []
|
|
self.tracking_bb = bb
|
|
self.tracking_mask = tracking_mask
|
|
|
|
corner_list = []
|
|
count = 1
|
|
while True:
|
|
print(f"Draw Corner #{count}")
|
|
matching_tpl_bb = cv2.selectROI("Matcher Reference", tracking_img, fromCenter=True, showCrosshair=True)
|
|
if bbox_center(matching_tpl_bb) == (0, 0):
|
|
break
|
|
|
|
matching_tpl_img = image_crop(image_processed_local, matching_tpl_bb)
|
|
self.corner_matcher_list.append(Corner(matching_tpl_img, matching_tpl_bb, name=f"Corner-{count}"))
|
|
corner_list.append(bbox_center(matching_tpl_bb))
|
|
|
|
self.matching_tpl_bb = matching_tpl_bb
|
|
|
|
print(f"Corner {count} added")
|
|
|
|
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
|
|
if self.params.pre_track:
|
|
self.tracker = cv2.TrackerKCF.create()
|
|
self.tracker.init(_image, self.tracking_bb)
|
|
|
|
# Refine initial corners and store them as reference
|
|
corners_local = self._corner_refine(image_processed_local, corners=corner_list)
|
|
|
|
# Transform corners to global
|
|
self.corner_ref = []
|
|
for corner_local in corners_local:
|
|
corner = (corner_local[0] + self.tracking_bb[0], corner_local[1] + self.tracking_bb[1])
|
|
self.corner_ref.append(corner)
|
|
|
|
return True
|
|
|
|
def process(self, _image: np.array, _image_anno: np.array):
|
|
# Update tracker bb
|
|
if self.tracker is not None:
|
|
_ok, self.tracking_bb = self.tracker.update(_image)
|
|
if not _ok:
|
|
return None
|
|
|
|
_image_local = image_crop(_image, self.tracking_bb)
|
|
CornerTracker.mask_apply(_image_local, self.tracking_mask)
|
|
_image_processed_local = CornerTracker.image_process(_image_local)
|
|
|
|
return self._match(_image_processed_local, _image_anno)
|
|
|
|
def _match(self, _image: np.array, _image_anno: np.array):
|
|
corner_list = []
|
|
for _ct in self.corner_matcher_list:
|
|
matcher_bbox_local, matcher = _ct.process(_image)
|
|
|
|
# Draw path
|
|
corner = bbox_center(matcher_bbox_local)
|
|
corner_list.append(corner)
|
|
|
|
self._debug(_image_anno, matcher_bbox_local)
|
|
|
|
# refine corners
|
|
corners_local = self._corner_refine(_image, corner_list)
|
|
|
|
# Transform corners to global
|
|
corners_refined = []
|
|
for corner_local in corners_local:
|
|
corner = (corner_local[0] + self.tracking_bb[0], corner_local[1] + self.tracking_bb[1])
|
|
corners_refined.append(corner)
|
|
|
|
corners_refined = np.array(corners_refined)
|
|
|
|
if self.params.show_path:
|
|
# Create path from global 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']), COLOR_TRACK, 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}] {corner_list[_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] + 0.5), int(corners_refined[_i, 1] + 0.5)), 4, COLOR_CIRCLE)
|
|
|
|
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__':
|
|
# Parse command line args
|
|
parser = argparse.ArgumentParser(description="Gearbox Tracker")
|
|
parser.add_argument("filename")
|
|
parser.add_argument("--scale", default=1.0)
|
|
parser.add_argument("--track", action="store_true")
|
|
parser.add_argument("--path", action="store_true")
|
|
parser.add_argument("--loop", action="store_true")
|
|
args = parser.parse_args()
|
|
video = cv2.VideoCapture(args.filename)
|
|
|
|
# Fill CornerTrackerParams from args
|
|
ct_params = CornerTrackerParams(_scale=np.float32(args.scale), _pre_track=args.track, _show_path=args.path)
|
|
|
|
# Let's go
|
|
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
|
|
tracker_count = 0
|
|
tracker_list = []
|
|
if not video.isOpened():
|
|
print('Cannot open video file')
|
|
sys.exit(1)
|
|
# Read first frame.
|
|
ok, image = video.read()
|
|
if not ok:
|
|
print('Cannot read video file')
|
|
sys.exit(1)
|
|
|
|
for tracker_count in range(0, 6):
|
|
select_window = image.copy()
|
|
print(f"Draw tracker #{tracker_count}")
|
|
ct = CornerTracker(ct_params, 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
|
|
|
|
# rewind to frame 0
|
|
video.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
|
|
|
key_wait = -1
|
|
dist_min = [(0, 0)]*len(tracker_list)
|
|
dist_max = [(0, 0)]*len(tracker_list)
|
|
while True:
|
|
# Start timer
|
|
timer = cv2.getTickCount()
|
|
# Read a new frame
|
|
ok, image = video.read()
|
|
if not ok:
|
|
if args.loop:
|
|
video.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
|
continue
|
|
else:
|
|
break
|
|
|
|
image_anno = image.copy()
|
|
i = 0
|
|
cv2.rectangle(image_anno, (25, 0), (int(image_anno.shape[1]), 25*(1+len(tracker_list))), (0, 0, 0), -1)
|
|
for ct in tracker_list:
|
|
mean_distance = ct.process(image, image_anno)
|
|
if mean_distance is not None:
|
|
scale = ct_params.scale
|
|
scaled_distance = (scale*mean_distance[0], scale*mean_distance[1])
|
|
dist_min[i] = (min(scaled_distance[0], dist_min[i][0]), min(scaled_distance[1], dist_min[i][1]))
|
|
dist_max[i] = (max(scaled_distance[0], dist_max[i][0]), max(scaled_distance[1], dist_max[i][1]))
|
|
cv2.putText(image_anno, f"Distance [{i}] : ({scaled_distance[0]:+05.2f}, "
|
|
f"{scaled_distance[1]:+05.2f}), Min: ({dist_min[i][0]:+05.2f}, {dist_min[i][1]:+05.2f}),"
|
|
f" Max: ({dist_max[i][0]:+05.2f}, {dist_max[i][1]:+05.2f})",
|
|
(25, 25 * (i + 1)), cv2.FONT_HERSHEY_SIMPLEX, 0.5, ct.color, 1)
|
|
i += 1
|
|
|
|
# Calculate Frames per second (FPS)
|
|
fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)
|
|
# Display FPS on frame
|
|
cv2.putText(image_anno, "FPS : " + str(int(fps)), (20, image_anno.shape[0] - 20),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (50, 170, 50), 1)
|
|
|
|
cv2.imshow(f"Image Anno", image_anno)
|
|
|
|
# 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
|
|
|