-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathvision_track_all.py
2709 lines (2118 loc) · 126 KB
/
vision_track_all.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# 系统题目：目标检测跟踪一体化系统
# 作者：ChenHu
# Copyright： Free
# 更新时间： 20230802
# 重点逻辑：
# 当界面涉及到不断更新（大量数据）的时候，会由于对象多次创建出现数据累积的情况。目前替换更新策略走不通，目前走的是销毁重建策略。
# 具体而说，文本更新用的是替换更新策略，图片更新用的是销毁重建策略（替换更新出现界面上无图的现象）
# 注意点：不更新的数据避免多次创建，小数据多次更新要留意积累现象（次要），大数据多次更新要留意积累现象（主要）
import tkinter as tk
import tkinter.ttk as ttk
import tkinter.messagebox
from PIL import Image, ImageTk
from yolo.utils.augmentations import (Albumentations, augment_hsv, classify_albumentations, classify_transforms, copy_paste,
letterbox, mixup, random_perspective)
import cv2
import random
from tkinter import filedialog, dialog
import torch
import matplotlib.pyplot as plt
import numpy as np
from tkinter import *
import datetime
import os
import time
import argparse
import os
import platform
import sys
from pathlib import Path
import tkinter.messagebox as messagebox
from PIL import Image, ImageTk, ImageDraw
import torch
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
from yolo.models.common import DetectMultiBackend
from yolo.utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams
from yolo.utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr,
cv2,
increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)
from yolo.utils.plots import Annotator, colors, save_one_box
from yolo.utils.torch_utils import select_device, smart_inference_mode
def load_model():
global tracker_siamcar
global tracker_siammask
global tracker_siamrpnpp
global hp
global opt_detect
global text_load_track,text_load_detect
# siammask
if siammask_choose.get():
import os
import argparse
import cv2
import torch
import numpy as np
from glob import glob
from siammask.pysot.core.config import cfg_mask
from siammask.pysot.models.model_builder import ModelBuilder
from siammask.pysot.tracker.tracker_builder import build_tracker
torch.set_num_threads(1)
parser = argparse.ArgumentParser(description='tracking demo')
parser.add_argument('--config', type=str, default="siammask/siammask_r50_l3/config.yaml",
help='config file')
parser.add_argument('--snapshot', type=str, default="siammask/siammask_r50_l3/model.pth",
help='model name')
parser.add_argument('--video_name', default='demo/bag.avi', type=str,
help='videos or image files')
args_mask = parser.parse_args()
cfg_mask.merge_from_file(args_mask.config)
cfg_mask.CUDA = torch.cuda.is_available() and cfg_mask.CUDA
cfg_mask.snapshot = args_mask.snapshot
device = torch.device('cuda' if cfg_mask.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(torch.load(args_mask.snapshot,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# build tracker
tracker_siammask = build_tracker(model)
# siamrpn++
if siamrpnpp_choose.get():
import os
import argparse
import cv2
import torch
import numpy as np
from glob import glob
from siamrpnpp.pysot.core.config import cfg as cfg_rpnpp
from siamrpnpp.pysot.models.model_builder import ModelBuilder
from siamrpnpp.pysot.tracker.tracker_builder import build_tracker
torch.set_num_threads(1)
parser = argparse.ArgumentParser(description='tracking demo')
parser.add_argument('--config', type=str, default="siamrpnpp/experiments/siamrpn_r50_l234_dwxcorr/config.yaml",
help='config file')
parser.add_argument('--snapshot', type=str, default="siamrpnpp/experiments/siamrpn_r50_l234_dwxcorr/model.pth",
help='model name')
parser.add_argument('--video_name', default='demo/bag.avi', type=str,
help='videos or image files')
args_rpnpp = parser.parse_args()
cfg_rpnpp.merge_from_file(args_rpnpp.config)
cfg_rpnpp.CUDA = torch.cuda.is_available() and cfg_rpnpp.CUDA
device = torch.device('cuda' if cfg_rpnpp.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model.load_state_dict(torch.load(args_rpnpp.snapshot,
map_location=lambda storage, loc: storage.cpu()))
model.eval().to(device)
# build tracker
tracker_siamrpnpp = build_tracker(model)
# siamcar
if siamcar_choose.get():
import os
import sys
sys.path.append('../')
import argparse
import cv2
import torch
from glob import glob
from siamcar.pysot.core.config import cfg
from siamcar.pysot.models.model_builder import ModelBuilder
from siamcar.pysot.tracker.siamcar_tracker import SiamCARTracker
from siamcar.pysot.utils.model_load import load_pretrain
torch.set_num_threads(1)
parser = argparse.ArgumentParser(description='SiamCAR demo')
parser.add_argument('--config', type=str, default='./siamcar/experiments/siamcar_r50/config.yaml',
help='config file')
parser.add_argument('--snapshot', type=str, default='./siamcar/snapshot/checkpoint_e20.pth', help='model name')
# parser.add_argument('--snapshot', type=str, default='./snapshot/SiamCAR-GOT.pth', help='model name')
parser.add_argument('--video_name', default='./Dancer', type=str, help='videos or image files')
args = parser.parse_args()
cfg.merge_from_file(args.config)
cfg.CUDA = torch.cuda.is_available()
device = torch.device('cuda' if cfg.CUDA else 'cpu')
# create model
model = ModelBuilder()
# load model
model = load_pretrain(model, args.snapshot).eval().to(device)
# build tracker
tracker_siamcar = SiamCARTracker(model, cfg.TRACK)
hp = {'lr': 0.3, 'penalty_k': 0.04, 'window_lr': 0.4}
# ----------------------------------------------
text_load_track = tkinter.Label(window, bd=10, font=("Microsoft YaHei", 15, "bold"), text="跟踪加载结束", fg="red")
text_load_track.place(relx=0.325, rely=0.90)
# 检测模型相关
opt_detect = parse_opt()
main_loadmodel(opt_detect)
text_load_detect = tkinter.Label(window, bd=10, font=("Microsoft YaHei", 15, "bold"), text="检测加载结束", fg="red")
text_load_detect.place(relx=0.325, rely=0.945)
def on_mouse_press(event):
global start_x, start_y,is_waiting_for_draw
start_x = event.x
start_y = event.y
def on_mouse_release(event,img_input,img_orishape):
global end_x, end_y, init_rect_draw
global start_x, start_y
global binding, binding_track, binding_track_release,is_waiting_for_draw
global tracker_siamcar,tracker_siamrpnpp,tracker_siammask
global cropped_img
end_x = event.x
end_y = event.y
# 创建一个临时画布
# temp_canvas = tk.Canvas(window, width=img_input.width, height=img_input.height)
# temp_canvas.place(relx=0.1, rely=0.1)
# 在临时画布上绘制矩形框
draw = ImageDraw.Draw(img_input)
draw.rectangle((start_x, start_y, end_x, end_y), outline="red", width=2)
# 将绘制后的图片转换为PhotoImage
photo_with_rect = ImageTk.PhotoImage(img_input)
# 在界面上更新图片
imglabel_track_frame.configure(image=photo_with_rect)
imglabel_track_frame.image = photo_with_rect
window.update()
# img_width = imglabel_track_frame.winfo_width() # 注意label的长宽和图片长宽并不一致，我们应该计算与图片的相对坐标
# img_height = imglabel_track_frame.winfo_height()
img_width = img_input.width # 显示图片本身长宽
img_height = img_input.height
im0_x = (start_x) * img_orishape.shape[1] / img_width
im0_y = (start_y) * img_orishape.shape[0] / img_height
# 计算点击事件相对于图片im0的坐标
im0_x_2 = (event.x) * img_orishape.shape[1] / img_width
im0_y_2 = (event.y) * img_orishape.shape[0] / img_height
cropped_img = img_orishape[int(im0_y):int(im0_y_2), int(im0_x):int(im0_x_2)]
init_rect_track = (im0_x, im0_y, im0_x_2 - im0_x, im0_y_2 - im0_y)
if siamcar_choose.get():
tracker_siamcar.init(img_orishape, init_rect_track) # 跟踪器初始化
if siammask_choose.get():
tracker_siammask.init(img_orishape, init_rect_track) # 跟踪器初始化
if siamrpnpp_choose.get():
tracker_siamrpnpp.init(img_orishape, init_rect_track) # 跟踪器初始化
# tracker.init(img_orishape, init_rect_track) # 跟踪器初始化
imglabel_track_frame.unbind("<ButtonPress-1>", binding_track)
imglabel_track_frame.unbind("<ButtonRelease-1>",binding_track_release)
print('绘制结束')
img_target = Image.fromarray(cropped_img[:,:,::-1])
img_target = image_resize(img_target, 100, 100)
img_target = ImageTk.PhotoImage(img_target) # 用PIL模块的PhotoImage打开
imglabel_target = tkinter.Label(window, bd=10, image=img_target)
imglabel_target.image = img_target
imglabel_target.place(relx=0.82, rely=0.13)
window.update()
is_waiting_for_draw = False
def handle_input(): # 输入帧数
global frame_choose
frame_choose = float(value_entry.get())
print("User input:", frame_choose)
frame_choose = frame_choose
#获得跟踪模型相关代码
import numpy as np
def det_and_track():
global mode_choose
mode_choose = "det_and_track"
text_det_and_track = tkinter.Label(window, bd=10, font=("Microsoft YaHei", 15, "bold"), text="一体模式", fg="red")
text_det_and_track.place(relx=0.045, rely=0.90)
def only_track():
global mode_choose
mode_choose = "only_track"
text_only_track = tkinter.Label(window, bd=10, font=("Microsoft YaHei", 15, "bold"), text="跟踪模式", fg="red")
text_only_track.place(relx=0.045, rely=0.90)
def crop_image_by_click(im0, xyxy_list, x, y):
global init_rect
# 将xyxy_list转换为NumPy数组
xyxy_array = np.array(xyxy_list)
# 检查每个目标框是否包含鼠标点击的位置
for xyxy in xyxy_array:
xmin, ymin, xmax, ymax = xyxy
# 如果点击位置在目标框内，则裁剪图像并返回
if x >= xmin and x <= xmax and y >= ymin and y <= ymax:
cropped_img = im0[int(ymin):int(ymax), int(xmin):int(xmax)]
init_rect = (xmin, ymin, xmax - xmin, ymax - ymin)
return cropped_img
# 如果点击位置不在任何目标框内，则返回提示信息
return "点击区域无有效目标"
def wait_for_mouse_click():
global is_waiting_for_click
global init_rect
is_waiting_for_click = True
while is_waiting_for_click:
window.update()
def wait_for_draw_target():
global is_waiting_for_draw
global init_rect
while is_waiting_for_draw:
window.update()
def handle_click(event,im0,xyxy_list):
from PIL import Image, ImageTk
global is_waiting_for_click
global clicked_x, clicked_y
global imglabel_track_frame
global imglabel_target
global img_target_flag
global track_flag
global template_img_choose
global init_rect
global prev_frame_time_2
if img_target_flag==0:
pass
else:
imglabel_target.destroy()
# 点击事件：根据点击位置最上层的窗口计算坐标
# # 计算点击事件相对于图片 im0 的坐标(相对总窗口window：未用)
# img_x = imglabel_track_frame.winfo_x() # 相对左上角坐标
# img_y = imglabel_track_frame.winfo_y()
# img_width = imglabel_track_frame.winfo_width() # 本身长宽
# img_height = imglabel_track_frame.winfo_height()
#
# # 计算点击事件相对于图片im0的坐标
# im0_x = (event.x - img_x) * im0.shape[1] / img_width
# im0_y = (event.y - img_y) * im0.shape[0] / img_height
#
# # 计算点击事件相对于图片 im0 的坐标(相对显式图片)
img_width = imglabel_track_frame.winfo_width() # 显示图片本身长宽
img_height = imglabel_track_frame.winfo_height()
#
# 计算点击事件相对于图片 im0 的坐标(相对显式图片)
# img_width = im0.widch # 显示图片本身长宽
# img_height = im0.height
# 计算点击事件相对于图片im0的坐标 (有偏差但可接受，想改掉偏差需要改为img_width为界面上图片大小，而不是界面窗口大小)
im0_x = (event.x) * im0.shape[1] / img_width
im0_y = (event.y) * im0.shape[0] / img_height
is_waiting_for_click = False
cropped_img = crop_image_by_click(im0, xyxy_list, im0_x, im0_y)
if isinstance(cropped_img, np.ndarray):
# 显示裁剪后的图像
img_target = Image.fromarray(cropped_img[:,:,::-1])
img_target = image_resize(img_target,100,100)
img_target = ImageTk.PhotoImage(img_target) # 用PIL模块的PhotoImage打开
imglabel_target = tkinter.Label(window, bd=10, image=img_target)
imglabel_target.image = img_target
imglabel_target.place(relx=0.82, rely=0.13)
window.update()
img_target_flag = 1
track_flag = 1
prev_frame_time_2 = time.time()
window.unbind("<Button-1>", binding)
template_img_choose = cropped_img
# tracker.init(im0, init_rect) # 跟踪器初始化
if siamcar_choose.get():
tracker_siamcar.init(im0, init_rect) # 跟踪器初始化
if siammask_choose.get():
tracker_siammask.init(im0, init_rect) # 跟踪器初始化
if siamrpnpp_choose.get():
tracker_siamrpnpp.init(im0, init_rect) # 跟踪器初始化
else:
if (event.x < img_width) and (event.y <img_height):
messagebox.showerror("错误", "未选中目标")
wait_for_mouse_click()
else:
messagebox.showerror("错误", "未启动跟踪模式")
window.unbind("<Button-1>", binding)
class VideoPlayer:
def __init__(self, master,video_path):
# 创建主窗口
self.master = master
# self.master.geometry('640x480')
# 创建标签用于显示视频帧
self.label = tk.Label(self.master)
self.label.pack()
# 创建播放按钮
self.play_button = tk.Button(self.master, text='Play', command=self.play_video)
self.play_button.pack(pady=10)
# 创建OpenCV视频对象
self.cap = cv2.VideoCapture(video_path)
def play_video(self,video_path):
# 循环读取视频帧并显示在标签中
while True:
ret, frame = self.cap.read()
if ret:
# 将OpenCV图像转换为PIL图像并显示在标签中
image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
image = Image.fromarray(image)
photo = ImageTk.PhotoImage(image)
self.label.configure(image=photo)
self.label.image = photo
# 等待一段时间
self.master.after(30, lambda: None)
else:
break
# 重置视频对象并重新播放
self.cap.release()
self.cap = cv2.VideoCapture(video_path)
def run(self):
self.master.mainloop()
def adjust_font_size(text):
text_width = text.winfo_width()
text_height = text.winfo_height()
text_size = int(text['font'].split()[1])
max_size = 50
min_size = 5
while (text.index('end-1c').split('.')[0] != 1 and
(text.winfo_width() < text_width or text.winfo_height() < text_height) and
text_size < max_size):
text_size += 1
text.configure(font=('Arial', text_size))
while text_size > min_size and (text.winfo_width() > text_width or text.winfo_height() > text_height):
text_size -= 1
text.configure(font=('Arial', text_size))
def image_resize(img, screen_width=800, screen_height=500):
image = img
raw_width, raw_height = image.size[0], image.size[1]
max_width, max_height = raw_width, screen_height
min_width = max(raw_width, max_width)
# 按照比例缩放
min_height = int(raw_height * min_width / raw_width)
# 第1次快速调整
while min_height > screen_height:
min_height = int(min_height * .9533)
# 第2次精确微调
while min_height < screen_height:
min_height += 1
# 按照比例缩放
min_width = int(raw_width * min_height / raw_height)
# 适应性调整
while min_width > screen_width:
min_width -= 1
# 按照比例缩放
min_height = int(raw_height * min_width / raw_width)
return image.resize((min_width, min_height))
def open_file_output():
from PIL import Image
'''
打开文件
:return:local_
'''
global file_path
global file_text
global photo
global img
global cap_flag
global first_choose
global text_load_track, text_load_detect
global track_flag
global mode_choose
if mode_choose != "det_and_track" and mode_choose != "only_track":
messagebox.showwarning("Warning", "未选择跟踪模式")
return
track_flag = 0
if first_choose == 1:
first_choose = 0
else:
text_begin.destroy()
text_end.destroy()
# text_save.destroy()
text_load_track.destroy()
text_load_detect.destroy()
cap_flag = 0
file_path = filedialog.askopenfilename(title=u'选择视频')
print('打开文件：', file_path)
if file_path is not None:
file_text = "文件路径为：" + file_path
cap = cv2.VideoCapture(file_path)
if not cap.isOpened():
print("Error opening video file")
ret, frame = cap.read()
rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
img = Image.fromarray(rgb_frame)
# 关闭视频文件
cap.release()
# img = Image.open(file_path) # 打开图片
# img = cv.imread(str(file_path))
img = image_resize(img)
photo = ImageTk.PhotoImage(img) # 用PIL模块的PhotoImage打开
imglabel = tkinter.Label(window, bd=10, image=photo)
imglabel.place(relx=0.1, rely=0.1)
def run_contral():
global opt_detect
main(opt_detect)
@smart_inference_mode()
def run_loadmodel(
weights=ROOT / 'yolo/yolov5s.pt', # model path or triton URL
source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam)
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
save_video = False,
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
vid_stride=1, # video frame-rate stride
):
from PIL import Image
global cap_flag
global text_begin
global text_end
global text_model_creat
global frame_choose
global track_flag
global tracker_siamcar,tracker_siamrpnpp,tracker_siammask
global hp
global model_detect
global seen_detect, windows_detect, dt_detect,dataset_detect,webcam_detect,save_dir_detect,vid_path_detect, vid_writer_detect,names_detect
track_flag = 0
if cap_flag==1:
source = str(source)
save_img = not nosave and not source.endswith('.txt') # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
webcam_detect = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
screenshot = source.lower().startswith('screen')
if is_url and is_file:
source = check_file(source) # download
# Directories
save_dir_detect = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
(save_dir_detect / 'labels' if save_txt else save_dir_detect).mkdir(parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model_detect = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, names_detect, pt = model_detect.stride, model_detect.names, model_detect.pt
imgsz = check_img_size(imgsz, s=stride) # check image size
# Dataloader
bs = 1 # batch_size
if webcam_detect:
view_img = check_imshow(warn=True)
dataset_detect = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
bs = len(dataset_detect)
elif screenshot:
dataset_detect = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
else:
# dataset_detect = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
dataset_detect = None
vid_path_detect, vid_writer_detect = [None] * bs, [None] * bs
# Run inference
model_detect.warmup(imgsz=(1 if pt or model_detect.triton else bs, 3, *imgsz)) # warmup
seen_detect, windows_detect, dt_detect = 0, [], (Profile(), Profile(), Profile())
else:
source = str(source)
save_img = not nosave and not source.endswith('.txt') # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
webcam_detect = source.isnumeric() or source.endswith('.streams') or (is_url and not is_file)
screenshot = source.lower().startswith('screen')
if is_url and is_file:
source = check_file(source) # download
# Directories
save_dir_detect = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run
(save_dir_detect / 'labels' if save_txt else save_dir_detect).mkdir(parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model_detect = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, names_detect, pt = model_detect.stride, model_detect.names, model_detect.pt
imgsz = check_img_size(imgsz, s=stride) # check image size
# Dataloader
bs = 1 # batch_size
if webcam_detect:
view_img = check_imshow(warn=True)
dataset_detect = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
bs = len(dataset_detect)
elif screenshot:
dataset_detect = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
else:
dataset_detect = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
vid_path_detect, vid_writer_detect = [None] * bs, [None] * bs
# Run inference
model_detect.warmup(imgsz=(1 if pt or model_detect.triton else bs, 3, *imgsz)) # warmup
seen_detect, windows_detect, dt_detect = 0, [], (Profile(), Profile(), Profile())
@smart_inference_mode()
def run(
weights=ROOT / 'yolo/yolov5s.pt', # model path or triton URL
source=ROOT / 'data/images', # file/dir/URL/glob/screen/0(webcam)
data=ROOT / 'data/coco128.yaml', # dataset.yaml path
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
save_video = False,
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
vid_stride=1, # video frame-rate stride
):
from PIL import Image
global cap_flag
global text_begin
global text_end
global text_model_creat
global frame_choose
global track_flag
global tracker_siamcar,tracker_siamrpnpp,tracker_siammask
global hp
global model_detect
global seen_detect, windows_detect, dt_detect,dataset_detect,webcam_detect,save_dir_detect,vid_path_detect, vid_writer_detect,names_detect
global prev_frame_time
global prev_frame_time_2
global imglabel_track_frame
global imglabel_dect_frame
global img_target_flag
global binding,binding_track,binding_track_release
global template_img_choose
global init_rect
global cap
import PIL
global position_fps_title
global position_fps
global is_waiting_for_draw
# 防止部件反复创建到界面（不会替换，而会累加，不同于普通变量）引起内存增加
position_label_flag = True
position_fps_flag = True
imglabel_tracker_flag = True
imglabel_dect_frame_flag = True
position_time_dect_flag = True
position_track_time_flag = True
global mode_choose
global init_rect_draw
if mode_choose != "det_and_track" and mode_choose != "only_track":
messagebox.showwarning("Warning", "未选择跟踪模式")
return
# 选择仅跟踪还是跟踪检测一体化
if mode_choose == "only_track" :
# 摄像头模式
if cap_flag == 1:
track_flag = 0
prev_frame_time_3 = time.time()
# text_model_creat.destroy()
text_begin = tkinter.Label(window, bd=10, font=("Microsoft YaHei", 15, "bold"), text="开始处理", fg="red")
text_begin.place(relx=0.465, rely=0.90)
window.update()
frame_count = 0
fps_num = 1
track_points_siamcar = []
track_points_siammask = []
track_points_siamrpnpp = []
# for path, im, im0s, vid_cap, s in dataset_detect: # 遍历视频序列的每一帧（直接遍历图片，不预先遍历序列）
while (cap.isOpened()):
# if fps_num % 10 == 0 and fps_num!=0 :
# del ret_flag,img,im0s,img_track,im,im0
# position_time_title.destroy()
# position_time.destroy()
# position_fps.destroy()
# position_fps_title.destroy()
fps_num = fps_num + 1
frame_count = frame_count + 1
ret_flag, img = cap.read()
img_track = img # 变色图片，但源代码用的是这个
im0s = img[:, :, ::-1] # 原图(恢复原色)
im = im0s
im0 = im0s.copy()
if show_track_3.get():
current_time = time.time()
frame_interval = current_time - prev_frame_time_3
# 更新上一帧时间戳为当前时间戳
# prev_frame_time = time.time()
prev_frame_time_3 = time.time()
if position_time_dect_flag:
position_time_title_dect = tkinter.Label(window, font=("Microsoft YaHei", 11, "bold"), text="")
position_time_title_dect.place(relx=0.781, rely=0.87)
position_time_dect = tkinter.Label(window, font=("Microsoft YaHei", 11, "bold"), text="")
position_time_dect.place(relx=0.81, rely=0.87)
position_time_title_dect.config(text=f"用时: ")
position_time_dect.config(text=f"{frame_interval:<.4f} s")
window.update()
position_time_dect_flag = False
# 更新标签的文本内容
else:
position_time_dect.config(text=f"{frame_interval:<.4f} s")
window.update()
if track_flag == 0:
im_vision = Image.fromarray(im0)
img = image_resize(im_vision)
photo = ImageTk.PhotoImage(img) # 用PIL模块的PhotoImage打开
imglabel_track_frame = tkinter.Label(window, bd=10, image=photo)
imglabel_track_frame.image = photo
imglabel_track_frame.place(relx=0.1, rely=0.1)
window.update()
if frame_count == int(frame_choose):
is_waiting_for_draw = True
track_flag = 1
print("绘制跟踪目标")
binding_track = imglabel_track_frame.bind("<ButtonPress-1>", on_mouse_press)
binding_track_release = imglabel_track_frame.bind("<ButtonRelease-1>", lambda event: on_mouse_release(event, img,img_track))
wait_for_draw_target()
if track_flag == 1:
# 计算跟踪时间
# if show_track_3.get():
# # current_time_3 = time.time()
# # frame_interval_2 = current_time_3 - prev_frame_time_3
# # # 更新上一帧时间戳为当前时间戳
# # prev_frame_time_3 = time.time()
#
# if position_track_time_flag:
# position_time_title = tkinter.Label(window, font=("Microsoft YaHei", 11, "bold"), text="")
# position_time_title.place(relx=0.781, rely=0.87)
# position_time = tkinter.Label(window, font=("Microsoft YaHei", 11, "bold"), text="")
# position_time.place(relx=0.81, rely=0.87)
# # 更新标签的文本内容
#
# position_time_title.config(text=f"用时: ")
# position_time.config(text=f"{frame_interval_2:<.4f} s")
# window.update()
#
# position_track_time_flag = False
#
# else:
# position_time_title.config(text=f"用时: ")
# position_time.config(text=f"{frame_interval_2:<.4f} s")
# window.update()
# template_img = template_img_choose
search_img = img_track
# 跟踪目标并获取边界框
if siamcar_choose.get():
outputs_siamcar = tracker_siamcar.track(search_img,hp)
bbox_siamcar = list(map(int, outputs_siamcar['bbox']))
cv2.rectangle(search_img, (bbox_siamcar[0], bbox_siamcar[1]),
(bbox_siamcar[0] + bbox_siamcar[2], bbox_siamcar[1] + bbox_siamcar[3]),
(0, 0, 255), 3)
# 添加目标轨迹点
track_points_siamcar.append((bbox_siamcar[0] + bbox_siamcar[2] // 2, bbox_siamcar[1] + bbox_siamcar[3] // 2))
if len(track_points_siamcar) > 80: # 轨迹长度
track_points_siamcar.pop(0)
# 绘制目标轨迹
if track_flag == 1 and show_track.get():
prev_point = None
for i, point in enumerate(track_points_siamcar):
size = int((i + 1) * 0.2)
if size >= 4:
size = 4
# color = (0, 255 - size * 20, 0)
color = (0, 0, 255)
# cv2.circle(search_img, point, size, color, -1)
# 连接前一个点和当前点
if prev_point is not None:
cv2.line(search_img, prev_point, point, color, thickness=2)
# 更新前一个点
prev_point = point
if siammask_choose.get():
outputs_siammask = tracker_siammask.track(search_img)
bbox_mask = list(map(int, outputs_siammask['bbox']))
cv2.rectangle(search_img, (bbox_mask[0], bbox_mask[1]),
(bbox_mask[0] + bbox_mask[2], bbox_mask[1] + bbox_mask[3]),
(0, 255, 0), 3)
# 添加目标轨迹点
track_points_siammask.append((bbox_mask[0] + bbox_mask[2] // 2, bbox_mask[1] + bbox_mask[3] // 2))
if len(track_points_siammask) > 80: # 轨迹长度
track_points_siammask.pop(0)
# 绘制目标轨迹
if track_flag == 1 and show_track.get():
prev_point = None
for i, point in enumerate(track_points_siammask):
size = int((i + 1) * 0.2)
if size >= 4:
size = 4
# color = (0, 255 - size * 20, 0)
color = (0, 255, 0)
# cv2.circle(search_img, point, size, color, -1)
if prev_point is not None:
cv2.line(search_img, prev_point, point, color, thickness=2)
# 更新前一个点
prev_point = point
if siamrpnpp_choose.get():
outputs_siamrpnpp = tracker_siamrpnpp.track(search_img)
bbox_rpnpp = list(map(int, outputs_siamrpnpp['bbox']))
cv2.rectangle(search_img, (bbox_rpnpp[0], bbox_rpnpp[1]),
(bbox_rpnpp[0] + bbox_rpnpp[2], bbox_rpnpp[1] + bbox_rpnpp[3]),
(255, 0, 0), 3)
# 添加目标轨迹点
track_points_siamrpnpp.append((bbox_rpnpp[0] + bbox_rpnpp[2] // 2, bbox_rpnpp[1] + bbox_rpnpp[3] // 2))
if len(track_points_siamrpnpp) > 80: # 轨迹长度
track_points_siamrpnpp.pop(0)
# 绘制目标轨迹
if track_flag == 1 and show_track.get():
prev_point = None
for i, point in enumerate(track_points_siamrpnpp):
size = int((i + 1) * 0.2)
if size >= 4:
size = 4
# color = (0, 255 - size * 20, 0)
color = (255, 0, 0)
# cv2.circle(search_img, point, size, color, -1)
if prev_point is not None:
cv2.line(search_img, prev_point, point, color, thickness=2)
# 更新前一个点
prev_point = point
# outputs = tracker.track(search_img)
# # 此处如果是原图则会出错，可理解为cv2默认格式是变色格式
# cv2.rectangle(search_img, (bbox[0], bbox[1]),
# (bbox[0] + bbox[2], bbox[1] + bbox[3]),
# (0, 255, 0), 3)
# # 添加目标轨迹点
# track_points.append((bbox[0] + bbox[2] // 2, bbox[1] + bbox[3] // 2))
# if len(track_points) > 80: # 轨迹长度
# track_points.pop(0)
#
# # 绘制目标轨迹
# if track_flag == 1 and show_track.get():