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