139 lines
4.0 KiB
Python
139 lines
4.0 KiB
Python
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], ")")
|
|
|
|
return corners_refined
|
|
|
|
|
|
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)
|
|
|
|
while True:
|
|
# Read a new frame
|
|
ok, image = video.read()
|
|
if not ok:
|
|
video.set(cv2.CAP_PROP_POS_FRAMES, 0)
|
|
continue
|
|
|
|
# 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, (9, 9), 0)
|
|
image_anno = image.copy()
|
|
|
|
if ok:
|
|
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, bbox_round(p['from']), bbox_round(p['to']), (0, 0, 255), 2)
|
|
|
|
# refine corners
|
|
corners_refined = refine_corner(cv2.cvtColor(image, cv2.COLOR_BGR2GRAY), np.array(corners, dtype=np.float32))
|
|
|
|
# show refined corners
|
|
for i in range(corners_refined.shape[0]):
|
|
print(" -- 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("Crop", crop)
|
|
cv2.imshow("Image", 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 = 30
|
|
else:
|
|
key_wait = -1
|