mirror of
https://github.com/opencv/opencv.git
synced 2024-11-24 03:00:14 +08:00
182 lines
8.4 KiB
Python
182 lines
8.4 KiB
Python
#!/usr/bin/env python
|
|
'''
|
|
Tracker demo
|
|
|
|
For usage download models by following links
|
|
For GOTURN:
|
|
goturn.prototxt and goturn.caffemodel: https://github.com/opencv/opencv_extra/tree/c4219d5eb3105ed8e634278fad312a1a8d2c182d/testdata/tracking
|
|
For DaSiamRPN:
|
|
network: https://www.dropbox.com/s/rr1lk9355vzolqv/dasiamrpn_model.onnx?dl=0
|
|
kernel_r1: https://www.dropbox.com/s/999cqx5zrfi7w4p/dasiamrpn_kernel_r1.onnx?dl=0
|
|
kernel_cls1: https://www.dropbox.com/s/qvmtszx5h339a0w/dasiamrpn_kernel_cls1.onnx?dl=0
|
|
For NanoTrack:
|
|
nanotrack_backbone: https://github.com/HonglinChu/SiamTrackers/blob/master/NanoTrack/models/nanotrackv2/nanotrack_backbone_sim.onnx
|
|
nanotrack_headneck: https://github.com/HonglinChu/SiamTrackers/blob/master/NanoTrack/models/nanotrackv2/nanotrack_head_sim.onnx
|
|
|
|
USAGE:
|
|
tracker.py [-h] [--input INPUT_VIDEO]
|
|
[--tracker_algo TRACKER_ALGO (mil, goturn, dasiamrpn, nanotrack, vittrack)]
|
|
[--goturn GOTURN_PROTOTXT]
|
|
[--goturn_model GOTURN_MODEL]
|
|
[--dasiamrpn_net DASIAMRPN_NET]
|
|
[--dasiamrpn_kernel_r1 DASIAMRPN_KERNEL_R1]
|
|
[--dasiamrpn_kernel_cls1 DASIAMRPN_KERNEL_CLS1]
|
|
[--nanotrack_backbone NANOTRACK_BACKBONE]
|
|
[--nanotrack_headneck NANOTRACK_TARGET]
|
|
[--vittrack_net VITTRACK_MODEL]
|
|
[--vittrack_net VITTRACK_MODEL]
|
|
[--tracking_score_threshold TRACKING SCORE THRESHOLD FOR ONLY VITTRACK]
|
|
[--backend CHOOSE ONE OF COMPUTATION BACKEND]
|
|
[--target CHOOSE ONE OF COMPUTATION TARGET]
|
|
'''
|
|
|
|
# Python 2/3 compatibility
|
|
from __future__ import print_function
|
|
|
|
import sys
|
|
|
|
import numpy as np
|
|
import cv2 as cv
|
|
import argparse
|
|
|
|
from video import create_capture, presets
|
|
|
|
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_HALIDE, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE, cv.dnn.DNN_BACKEND_OPENCV,
|
|
cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
|
|
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD,
|
|
cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)
|
|
|
|
class App(object):
|
|
|
|
def __init__(self, args):
|
|
self.args = args
|
|
self.trackerAlgorithm = args.tracker_algo
|
|
self.tracker = self.createTracker()
|
|
|
|
def createTracker(self):
|
|
if self.trackerAlgorithm == 'mil':
|
|
tracker = cv.TrackerMIL_create()
|
|
elif self.trackerAlgorithm == 'goturn':
|
|
params = cv.TrackerGOTURN_Params()
|
|
params.modelTxt = self.args.goturn
|
|
params.modelBin = self.args.goturn_model
|
|
tracker = cv.TrackerGOTURN_create(params)
|
|
elif self.trackerAlgorithm == 'dasiamrpn':
|
|
params = cv.TrackerDaSiamRPN_Params()
|
|
params.model = self.args.dasiamrpn_net
|
|
params.kernel_cls1 = self.args.dasiamrpn_kernel_cls1
|
|
params.kernel_r1 = self.args.dasiamrpn_kernel_r1
|
|
params.backend = args.backend
|
|
params.target = args.target
|
|
tracker = cv.TrackerDaSiamRPN_create(params)
|
|
elif self.trackerAlgorithm == 'nanotrack':
|
|
params = cv.TrackerNano_Params()
|
|
params.backbone = args.nanotrack_backbone
|
|
params.neckhead = args.nanotrack_headneck
|
|
params.backend = args.backend
|
|
params.target = args.target
|
|
tracker = cv.TrackerNano_create(params)
|
|
elif self.trackerAlgorithm == 'vittrack':
|
|
params = cv.TrackerVit_Params()
|
|
params.net = args.vittrack_net
|
|
params.tracking_score_threshold = args.tracking_score_threshold
|
|
params.backend = args.backend
|
|
params.target = args.target
|
|
tracker = cv.TrackerVit_create(params)
|
|
else:
|
|
sys.exit("Tracker {} is not recognized. Please use one of three available: mil, goturn, dasiamrpn, nanotrack.".format(self.trackerAlgorithm))
|
|
return tracker
|
|
|
|
def initializeTracker(self, image):
|
|
while True:
|
|
print('==> Select object ROI for tracker ...')
|
|
bbox = cv.selectROI('tracking', image)
|
|
print('ROI: {}'.format(bbox))
|
|
if bbox[2] <= 0 or bbox[3] <= 0:
|
|
sys.exit("ROI selection cancelled. Exiting...")
|
|
|
|
try:
|
|
self.tracker.init(image, bbox)
|
|
except Exception as e:
|
|
print('Unable to initialize tracker with requested bounding box. Is there any object?')
|
|
print(e)
|
|
print('Try again ...')
|
|
continue
|
|
|
|
return
|
|
|
|
def run(self):
|
|
videoPath = self.args.input
|
|
print('Using video: {}'.format(videoPath))
|
|
camera = create_capture(cv.samples.findFileOrKeep(videoPath), presets['cube'])
|
|
if not camera.isOpened():
|
|
sys.exit("Can't open video stream: {}".format(videoPath))
|
|
|
|
ok, image = camera.read()
|
|
if not ok:
|
|
sys.exit("Can't read first frame")
|
|
assert image is not None
|
|
|
|
cv.namedWindow('tracking')
|
|
self.initializeTracker(image)
|
|
|
|
print("==> Tracking is started. Press 'SPACE' to re-initialize tracker or 'ESC' for exit...")
|
|
|
|
while camera.isOpened():
|
|
ok, image = camera.read()
|
|
if not ok:
|
|
print("Can't read frame")
|
|
break
|
|
|
|
ok, newbox = self.tracker.update(image)
|
|
#print(ok, newbox)
|
|
|
|
if ok:
|
|
cv.rectangle(image, newbox, (200,0,0))
|
|
|
|
cv.imshow("tracking", image)
|
|
k = cv.waitKey(1)
|
|
if k == 32: # SPACE
|
|
self.initializeTracker(image)
|
|
if k == 27: # ESC
|
|
break
|
|
|
|
print('Done')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print(__doc__)
|
|
parser = argparse.ArgumentParser(description="Run tracker")
|
|
parser.add_argument("--input", type=str, default="vtest.avi", help="Path to video source")
|
|
parser.add_argument("--tracker_algo", type=str, default="nanotrack", help="One of available tracking algorithms: mil, goturn, dasiamrpn, nanotrack, vittrack")
|
|
parser.add_argument("--goturn", type=str, default="goturn.prototxt", help="Path to GOTURN architecture")
|
|
parser.add_argument("--goturn_model", type=str, default="goturn.caffemodel", help="Path to GOTERN model")
|
|
parser.add_argument("--dasiamrpn_net", type=str, default="dasiamrpn_model.onnx", help="Path to onnx model of DaSiamRPN net")
|
|
parser.add_argument("--dasiamrpn_kernel_r1", type=str, default="dasiamrpn_kernel_r1.onnx", help="Path to onnx model of DaSiamRPN kernel_r1")
|
|
parser.add_argument("--dasiamrpn_kernel_cls1", type=str, default="dasiamrpn_kernel_cls1.onnx", help="Path to onnx model of DaSiamRPN kernel_cls1")
|
|
parser.add_argument("--nanotrack_backbone", type=str, default="nanotrack_backbone_sim.onnx", help="Path to onnx model of NanoTrack backBone")
|
|
parser.add_argument("--nanotrack_headneck", type=str, default="nanotrack_head_sim.onnx", help="Path to onnx model of NanoTrack headNeck")
|
|
parser.add_argument("--vittrack_net", type=str, default="vitTracker.onnx", help="Path to onnx model of vittrack")
|
|
parser.add_argument('--tracking_score_threshold', type=float, help="Tracking score threshold. If a bbox of score >= 0.3, it is considered as found ")
|
|
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
|
|
help="Choose one of computation backends: "
|
|
"%d: automatically (by default), "
|
|
"%d: Halide language (http://halide-lang.org/), "
|
|
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
|
|
"%d: OpenCV implementation, "
|
|
"%d: VKCOM, "
|
|
"%d: CUDA"% backends)
|
|
parser.add_argument("--target", choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
|
|
help="Choose one of target computation devices: "
|
|
'%d: CPU target (by default), '
|
|
'%d: OpenCL, '
|
|
'%d: OpenCL fp16 (half-float precision), '
|
|
'%d: VPU, '
|
|
'%d: VULKAN, '
|
|
'%d: CUDA, '
|
|
'%d: CUDA fp16 (half-float preprocess)'% targets)
|
|
|
|
args = parser.parse_args()
|
|
App(args).run()
|
|
cv.destroyAllWindows()
|