159 lines
3.9 KiB
Python
159 lines
3.9 KiB
Python
import cv2
|
|
import sys
|
|
import numpy as np
|
|
import imutils as im
|
|
|
|
IMG_SCALE_UP = 1
|
|
SEARCH_AREA = 0.2
|
|
|
|
(major_ver, minor_ver, subminor_ver) = cv2.__version__.split('.')
|
|
print(cv2.__version__)
|
|
|
|
def bbox_extend(bbox: cv2.typing.Rect, search_area: float):
|
|
we2 = int(round(bbox[2]*search_area/2))
|
|
he2 = int(round(bbox[3]*search_area/2))
|
|
res = (int(round(bbox[0]-we2)),)
|
|
res += (int(round(bbox[1]-he2)),)
|
|
res += (bbox[2]+2*we2,)
|
|
res += (bbox[3]+2*he2,)
|
|
|
|
return res
|
|
|
|
|
|
def image_crop(src, bbox: cv2.typing.Rect, search_area: float = 0):
|
|
bbox = bbox_extend(bbox, search_area)
|
|
x = bbox[0]
|
|
y = bbox[1]
|
|
w = bbox[2]
|
|
h = bbox[3]
|
|
|
|
dst = src[y:y+h, x:x+w]
|
|
return dst
|
|
|
|
|
|
def process_image_bbox(src):
|
|
# convert to gray
|
|
gray = cv2.cvtColor(src, cv2.COLOR_BGR2GRAY)
|
|
|
|
return gray
|
|
|
|
|
|
def match(image_bbox, image_template, current_angle):
|
|
best_angle = 0
|
|
best_r = 0
|
|
best_maxLoc = 0
|
|
|
|
for dangle in np.linspace(-1.0, +1.0, 10):
|
|
# scale up to IMG_SCALE_UP-size
|
|
image_template_rotated = im.resize(image_template, IMG_SCALE_UP*w, IMG_SCALE_UP*h)
|
|
# rotated
|
|
image_template_rotated = im.rotate(image_template_rotated, current_angle + dangle)
|
|
# scale to original-size
|
|
image_template_rotated = im.resize(image_template_rotated, w, h)
|
|
|
|
# Perform template match operations
|
|
res = cv2.matchTemplate(image_bbox, image_template_rotated, cv2.TM_CCOEFF_NORMED)
|
|
(minVal, maxVal, minLoc, maxLoc) = cv2.minMaxLoc(res)
|
|
|
|
r = maxVal
|
|
if r > best_r:
|
|
best_r = r
|
|
best_angle = dangle
|
|
best_maxLoc = maxLoc
|
|
|
|
current_angle += best_angle
|
|
|
|
return current_angle, best_maxLoc
|
|
|
|
|
|
trackers = {
|
|
'KCF': cv2.TrackerKCF,
|
|
'CSRT': cv2.TrackerCSRT,
|
|
'MIL': cv2.TrackerMIL,
|
|
'GOTURN': cv2.TrackerGOTURN,
|
|
'NANO': cv2.TrackerNano,
|
|
'MOSSE': cv2.legacy.TrackerMOSSE,
|
|
'TLD': cv2.legacy.TrackerTLD,
|
|
'BOOSTING': cv2.legacy.TrackerBoosting,
|
|
'MEDIANFLOW': cv2.legacy.TrackerMedianFlow
|
|
}
|
|
tracker_type = 'KCF'
|
|
tracker = trackers[tracker_type].create()
|
|
|
|
video = cv2.VideoCapture("./data/production_id 4525346 (1080p).mp4")
|
|
|
|
# Exit if video not open
|
|
# d.
|
|
if not video.isOpened():
|
|
print("Could not open video")
|
|
sys.exit()
|
|
|
|
# Read first frame.
|
|
ok, frame = video.read()
|
|
if not ok:
|
|
print('Cannot read video file')
|
|
sys.exit()
|
|
|
|
# Uncomment the line below to select a different bounding box
|
|
bbox = cv2.selectROI("Frame", frame, False)
|
|
image_template = process_image_bbox(image_crop(frame, bbox, 0.5))
|
|
w, h = image_template.shape[::-1]
|
|
|
|
# Initialize tracker with first frame and bounding box
|
|
tracker.create()
|
|
tracker.init(frame, bbox)
|
|
|
|
current_angle = 0
|
|
while True:
|
|
# Read a new frame
|
|
ok, frame = video.read()
|
|
if not ok:
|
|
break
|
|
|
|
# Start timer
|
|
timer = cv2.getTickCount()
|
|
|
|
# Update tracker
|
|
ok, bbox = tracker.update(frame)
|
|
|
|
image_bbox = process_image_bbox(image_crop(frame, bbox, SEARCH_AREA))
|
|
image_template_rotated = image_template
|
|
|
|
if ok:
|
|
current_angle, best_maxLoc = match(image_bbox, image_template, current_angle)
|
|
image_template_rotated = im.rotate(image_template, current_angle)
|
|
|
|
# determine the starting and ending (x, y)-coordinates of the
|
|
# bounding box
|
|
(startX, startY) = best_maxLoc
|
|
startX += bbox[0]
|
|
startY += bbox[1]
|
|
endX = startX + bbox[2]
|
|
endY = startY + bbox[3]
|
|
|
|
# draw the bounding box on the image
|
|
cv2.rectangle(frame, (startX, startY), (endX, endY), (255, 0, 0), 3)
|
|
|
|
# Calculate Frames per second (FPS)
|
|
fps = cv2.getTickFrequency() / (cv2.getTickCount() - timer)
|
|
else:
|
|
# Tracking failure
|
|
cv2.putText(frame, "Tracking failure detected", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2)
|
|
|
|
# Display tracker type on frame
|
|
cv2.putText(frame, tracker_type + " Tracker", (100, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
|
|
|
|
# Display FPS on frame
|
|
cv2.putText(frame, "FPS : " + str(int(fps)), (100, 50), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (50, 170, 50), 2)
|
|
|
|
# Display result
|
|
cv2.imshow("Tracking", frame)
|
|
cv2.imshow("image_bbox", image_template_rotated)
|
|
|
|
# Exit if ESC pressed
|
|
k = cv2.waitKey(1) & 0xff
|
|
if k == 27:
|
|
break
|
|
|
|
|