1+ #!/usr/bin/env python
2+ # -*- coding: utf-8 -*-
3+ import re
4+ import sys
5+ import copy
6+ import time
7+ import argparse
8+ 9+ import cv2 as cv
10+ 11+ 12+ def get_args ():
13+ parser = argparse .ArgumentParser ()
14+ 15+ parser .add_argument ("--device" , default = "sample_movie/bird.mp4" )
16+ parser .add_argument ("--width" , help = 'cap width' , type = int , default = 960 )
17+ parser .add_argument ("--height" , help = 'cap height' , type = int , default = 540 )
18+ 19+ args = parser .parse_args ()
20+ 21+ return args
22+ 23+ 24+ def isint (s ):
25+ p = '[-+]?\d+'
26+ return True if re .fullmatch (p , s ) else False
27+ 28+ 29+ def initialize_tracker (window_name , image ):
30+ params = cv .TrackerDaSiamRPN_Params ()
31+ params .model = "model/DaSiamRPN/dasiamrpn_model.onnx"
32+ params .kernel_r1 = "model/DaSiamRPN/dasiamrpn_kernel_r1.onnx"
33+ params .kernel_cls1 = "model/DaSiamRPN/dasiamrpn_kernel_cls1.onnx"
34+ tracker = cv .TrackerDaSiamRPN_create (params )
35+ 36+ # 追跡対象指定
37+ while True :
38+ bbox = cv .selectROI (window_name , image )
39+ 40+ try :
41+ tracker .init (image , bbox )
42+ except Exception as e :
43+ print (e )
44+ continue
45+ 46+ return tracker
47+ 48+ 49+ def main ():
50+ color_list = [
51+ [255 , 0 , 0 ], # blue
52+ ]
53+ 54+ # 引数解析 #################################################################
55+ args = get_args ()
56+ 57+ cap_device = args .device
58+ cap_width = args .width
59+ cap_height = args .height
60+ 61+ # カメラ準備 ###############################################################
62+ if isint (cap_device ):
63+ cap_device = int (cap_device )
64+ cap = cv .VideoCapture (cap_device )
65+ cap .set (cv .CAP_PROP_FRAME_WIDTH , cap_width )
66+ cap .set (cv .CAP_PROP_FRAME_HEIGHT , cap_height )
67+ 68+ # Tracker初期化 ############################################################
69+ window_name = 'Tracker Demo'
70+ cv .namedWindow (window_name )
71+ 72+ ret , image = cap .read ()
73+ if not ret :
74+ sys .exit ("Can't read first frame" )
75+ tracker = initialize_tracker (window_name , image )
76+ 77+ while cap .isOpened ():
78+ ret , image = cap .read ()
79+ if not ret :
80+ break
81+ debug_image = copy .deepcopy (image )
82+ 83+ # 追跡アップデート
84+ start_time = time .time ()
85+ ok , bbox = tracker .update (image )
86+ elapsed_time = time .time () - start_time
87+ if ok :
88+ # 追跡後のバウンディングボックス描画
89+ cv .rectangle (debug_image , bbox , color_list [0 ], thickness = 2 )
90+ 91+ # 各アルゴリズム処理時間描画
92+ cv .putText (
93+ debug_image ,
94+ 'DaSiamRPN' + " : " + '{:.1f}' .format (elapsed_time * 1000 ) + "ms" ,
95+ (10 , 25 ), cv .FONT_HERSHEY_SIMPLEX , 0.7 , color_list [0 ], 2 ,
96+ cv .LINE_AA )
97+ 98+ cv .imshow (window_name , debug_image )
99+ 100+ k = cv .waitKey (1 )
101+ if k == 32 : # SPACE
102+ # 追跡対象再指定
103+ tracker = initialize_tracker (window_name , image )
104+ if k == 27 : # ESC
105+ break
106+ 107+ 108+ if __name__ == '__main__' :
109+ main ()
0 commit comments