540 lines
17 KiB
Python
540 lines
17 KiB
Python
import cv2
|
|
import numpy as np
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import csv
|
|
import json
|
|
import matplotlib.pyplot as plt
|
|
|
|
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
|
|
|
|
def from_dict(self, params: dict):
|
|
self.__dict__.update(params)
|
|
return self
|
|
|
|
|
|
class CornerTrackerSettings:
|
|
def __init__(self, _id: int, _mask_bb=None, _tracking_bb=None):
|
|
self.id = _id
|
|
self.mask_bb = _mask_bb
|
|
self.tracking_bb = _tracking_bb
|
|
self.matching_tpl_bb = []
|
|
|
|
def from_dict(self, _settings: dict):
|
|
self.__dict__.update(_settings)
|
|
return self
|
|
|
|
def append(self, _bb: np.array):
|
|
self.matching_tpl_bb.append(_bb)
|
|
|
|
|
|
class CornerTracker:
|
|
def __init__(self, _id: int, _params: CornerTrackerParams, color=(0, 255, 0), name: str = 'CornerTracker'):
|
|
self.id = _id
|
|
self.params = _params
|
|
self.color = color
|
|
self.name = name
|
|
self.settings = 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)
|
|
|
|
cv2.putText(_image_anno, f"{self.id}", (self.tracking_bb[0], self.tracking_bb[1] - 4),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.8, self.color, 2)
|
|
|
|
@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])
|
|
|
|
return bb
|
|
|
|
@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)
|
|
'''
|
|
_result = cv2.GaussianBlur(_result, (9, 9), 0)
|
|
_result = cv2.medianBlur(_result, 3)
|
|
'''
|
|
return _result
|
|
|
|
def create_from_settings(self, _image: np.array, _settings: CornerTrackerSettings, _image_anno: np.array):
|
|
_mask = CornerTracker.create_mask(_image, _settings.mask_bb)
|
|
image_processed_local, _ = self.create_tracking(_image, _mask, _settings.tracking_bb)
|
|
|
|
return self.create_matching(image_processed_local, _settings.tracking_bb, _settings.matching_tpl_bb, _image_anno)
|
|
|
|
@staticmethod
|
|
def create_mask(_image: np.array, mask_bb: np.array):
|
|
# Init mask
|
|
_mask = CornerTracker.mask_init(_image)
|
|
rect_upper_left = (mask_bb[0], mask_bb[1])
|
|
rect_lower_right = (mask_bb[0] + mask_bb[2], mask_bb[1] + mask_bb[3])
|
|
# Mask cut out
|
|
cv2.rectangle(_mask, rect_upper_left, rect_lower_right, 255, -1)
|
|
|
|
return _mask
|
|
|
|
def create_tracking(self, _image: np.array, _mask: np.array, tracking_bb: np.array):
|
|
# Init search area
|
|
tracking_img = image_crop(_image.copy(), tracking_bb)
|
|
tracking_mask = image_crop(_mask, tracking_bb)
|
|
CornerTracker.mask_apply(tracking_img, tracking_mask)
|
|
image_processed_local = CornerTracker.image_process(tracking_img)
|
|
|
|
# Initialize tracker with first frame and bounding box
|
|
if self.params.pre_track:
|
|
self.tracker = cv2.TrackerKCF.create()
|
|
self.tracker.init(_image, tracking_bb)
|
|
|
|
self.tracking_bb = tracking_bb
|
|
self.tracking_mask = image_crop(_mask, tracking_bb)
|
|
|
|
return image_processed_local, tracking_img
|
|
|
|
def create_matching(self, image_processed_local: np.array, tracking_bb: np.array, matching_tpl_bb_list: list, _image_anno: np.array):
|
|
# Create corner matcher
|
|
self.corner_matcher_list = []
|
|
corner_list = []
|
|
count = 1
|
|
for matching_tpl_bb in matching_tpl_bb_list:
|
|
if bbox_center(matching_tpl_bb) == (0, 0):
|
|
break
|
|
corner_list.append(bbox_center(matching_tpl_bb))
|
|
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}"))
|
|
|
|
print(f"Corner {count} added")
|
|
count += 1
|
|
|
|
if len(corner_list) == 0:
|
|
return False
|
|
|
|
# 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] + tracking_bb[0], corner_local[1] + tracking_bb[1])
|
|
self.corner_ref.append(corner)
|
|
|
|
for matching_tpl_bb in matching_tpl_bb_list:
|
|
self._debug(_image_anno, matching_tpl_bb)
|
|
|
|
return True
|
|
|
|
def init_reference_frame(self, _image: np.array, _image_anno: np.array):
|
|
_settings = CornerTrackerSettings(self.id)
|
|
# Draw mask
|
|
_mask = CornerTracker.mask_init(_image)
|
|
|
|
# Init mask
|
|
print(f"Draw mask object")
|
|
_mask_bb = cv2.selectROI("Tracker Reference", _image_anno, False)
|
|
cv2.destroyWindow("Tracker Reference")
|
|
if bbox_center(_mask_bb) == (0, 0):
|
|
_mask_bb = (0, 0, _image.shape[1], _image.shape[0])
|
|
|
|
_settings.mask_bb = _mask_bb
|
|
_mask = self.create_mask(_image, _mask_bb)
|
|
|
|
# Select search area
|
|
print(f"Select tracking object")
|
|
tracking_bb = cv2.selectROI("Tracker Reference", _image_anno, False)
|
|
cv2.destroyWindow("Tracker Reference")
|
|
if bbox_center(tracking_bb) == (0, 0):
|
|
return None
|
|
|
|
_settings.tracking_bb = tracking_bb
|
|
image_processed_local, tracking_img = self.create_tracking(_image, _mask, tracking_bb)
|
|
|
|
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
|
|
|
|
_settings.matching_tpl_bb.append(matching_tpl_bb)
|
|
print(f"Corner {count} added")
|
|
count += 1
|
|
|
|
cv2.destroyWindow("Matcher Reference")
|
|
print(f"Added {count} corners")
|
|
|
|
if self.create_matching(image_processed_local, _settings.tracking_bb, _settings.matching_tpl_bb, _image_anno):
|
|
return _settings
|
|
|
|
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)
|
|
|
|
|
|
def project_load(path: str = "./"):
|
|
_settings = None
|
|
try:
|
|
with open(os.path.join(path, "ocv_tracker_prj.json"), "r") as fp:
|
|
_settings = json.load(fp)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
return _settings
|
|
|
|
|
|
def project_save(_settings: dict, path: str = "./"):
|
|
with open(os.path.join(path, "ocv_tracker_prj.json"), "w") as fp:
|
|
json.dump(_settings, fp, indent=4)
|
|
|
|
|
|
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")
|
|
parser.add_argument("--csv", action="store_true")
|
|
parser.add_argument("--plot", action="store_true")
|
|
parser.add_argument("--name", default="Default")
|
|
args = parser.parse_args()
|
|
video = cv2.VideoCapture(args.filename)
|
|
|
|
# Fill CornerTrackerParams from args
|
|
ct_params = CornerTrackerParams(_scale=args.scale, _pre_track=args.track, _show_path=args.path)
|
|
|
|
prj = project_load(os.path.dirname(args.filename))
|
|
prj_loaded = False
|
|
if prj is not None:
|
|
ct_params = ct_params.from_dict(prj['params'])
|
|
prj_loaded = True
|
|
|
|
# Let's go
|
|
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)
|
|
|
|
colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 0, 255), (0, 255, 255), (255, 255, 255)]
|
|
tracker_count = 0
|
|
tracker_list = []
|
|
image_anno = image.copy()
|
|
|
|
if prj_loaded:
|
|
ct_params = CornerTrackerParams().from_dict(prj['params'])
|
|
for tracker_settings in prj['trackers']:
|
|
print(tracker_settings)
|
|
ct = CornerTracker(tracker_settings['id'], ct_params, colors[tracker_count], name=f"Tracker-{tracker_settings['id']}")
|
|
ct.create_from_settings(image, CornerTrackerSettings(tracker_count).from_dict(tracker_settings), image_anno)
|
|
tracker_list.append(ct)
|
|
tracker_count += 1
|
|
|
|
else:
|
|
prj = {'params': ct_params.__dict__, 'trackers': []}
|
|
|
|
cv2.imshow("Tracker Reference", image_anno)
|
|
|
|
print("Press 'a' for adding trackers, press 'd' to delete a tracker, press other key to continue")
|
|
while True:
|
|
print(f"Add tracker #{tracker_count}")
|
|
ct = CornerTracker(tracker_count, ct_params, colors[tracker_count], name=f"Tracker-{tracker_count}")
|
|
settings = ct.init_reference_frame(image, image_anno)
|
|
if settings is not None:
|
|
prj['trackers'].append(settings.__dict__)
|
|
tracker_list.append(ct)
|
|
print(f"Number of active tracker: {len(tracker_list)}")
|
|
else:
|
|
break
|
|
tracker_count += 1
|
|
|
|
project_save(prj, os.path.dirname(args.filename))
|
|
|
|
# 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)
|
|
result_csv = [[] for i in range(len(tracker_list))]
|
|
result_plot = [[] for i in range(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 = np.float32(ct_params.scale)
|
|
scaled_distance = (scale*mean_distance[0], scale*mean_distance[1])
|
|
if args.csv:
|
|
result_csv[i].append({f'dx': scaled_distance[0], 'dy': scaled_distance[1]})
|
|
if args.plot:
|
|
result_plot[i].append(scaled_distance)
|
|
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 [{ct.id}] : ({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
|
|
|
|
# writing to csv file
|
|
if len(result_csv[0]) > 0:
|
|
i = 1
|
|
for result in result_csv:
|
|
filename = f"{args.name}_{i}.csv"
|
|
with open(filename, 'w') as csvfile:
|
|
# creating a csv dict writer object
|
|
writer = csv.DictWriter(csvfile, fieldnames=['dx', 'dy'])
|
|
|
|
# writing headers (field names)
|
|
writer.writeheader()
|
|
|
|
# writing data rows
|
|
writer.writerows(result)
|
|
print(f"Exported {filename}")
|
|
i += 1
|
|
|
|
if len(result_plot[0]) > 0:
|
|
N_p = len(result_plot)
|
|
i = 1
|
|
for result in result_plot:
|
|
Np = len
|
|
n = range(0, len(result))
|
|
plt.subplot(N_p, 1, i)
|
|
plt.plot(n, np.array(result)[:, 0], n, np.array(result)[:, 1])
|
|
plt.grid()
|
|
plt.legend([f"x", f"y"])
|
|
i += 1
|
|
|
|
plt.show()
|