1 import mmcv
 2 import tempfile
 3 from mmtrack.apis import inference_sot, init_model
 4 
 5 # 输入输出视频路径
 6 input_video = 'data/bee.mp4'
 7 output = 'outputs/C4_SOT_bee_trace.mp4'
 8 
 9 # 指定单目标追踪算法 config 配置文件
10 sot_config = './configs/sot/siamese_rpn/siamese_rpn_r50_20e_lasot.py'
11 # 指定单目标检测算法的模型权重文件
12 sot_checkpoint = 'https://download.openmmlab.com/mmtracking/sot/siamese_rpn/siamese_rpn_r50_1x_lasot/siamese_rpn_r50_1x_lasot_20211203_151612-da4b3c66.pth'
13 # 初始化单目标追踪模型
14 sot_model = init_model(sot_config, sot_checkpoint, device='cuda:0')
15 
16 # 指定初始目标矩形框坐标 [x, y, w, h]
17 init_bbox_xywh = [132, 59, 57, 61]
18 
19 # 转成 [x1, y1, x2, y2 ]
20 init_bbox_xyxy = [init_bbox_xywh[0], init_bbox_xywh[1], init_bbox_xywh[0]+init_bbox_xywh[2], init_bbox_xywh[1]+init_bbox_xywh[3]]
21 
22 # 读入待预测视频
23 imgs = mmcv.VideoReader(input_video)
24 prog_bar = mmcv.ProgressBar(len(imgs))
25 out_dir = tempfile.TemporaryDirectory()
26 out_path = out_dir.name
27 
28 # 逐帧输入模型预测
29 circle_coord_list = []
30 print('开始逐帧处理')
31 for i, img in enumerate(imgs):
32     img_draw = img.copy()
33     
34     result = inference_sot(sot_model, img, init_bbox_xyxy, frame_id=i)
35     # 目标检测矩形框坐标
36     result_bbox = result['track_bboxes'][:4].astype('uint32')
37     
38     # 绘制目标检测矩形框：图像，左上角坐标，右下角坐标，颜色，线宽
39     img_draw = cv2.rectangle(img_draw, (result_bbox[0], result_bbox[1]), (result_bbox[2], result_bbox[3]), (0,255,0), 2)    
40     
41     # 获取矩形框中心点轨迹点坐标
42     circle_x = int((result_bbox[0] + result_bbox[2]) / 2)
43     circle_y = int((result_bbox[1] + result_bbox[3]) / 2)
44     circle_coord_list.append([circle_x, circle_y])
45     # 绘制从第一帧到当前帧的轨迹
46     for each in circle_coord_list:
47         # 绘制圆，指定圆心坐标和半径，红色，最后一个参数为线宽，-1表示填充
48         img_draw = cv2.circle(img_draw, (each[0],each[1]), 5, (0,0,255), -1)
49     
50     # 将当前帧的可视化效果保存为图片文件
51     cv2.imwrite(f'{out_path}/{i:06d}.jpg', img_draw)
52     prog_bar.update()
53 
54 # 将保存下来的各帧图片文件串成视频
55 print('导出视频，FPS {}'.format(imgs.fps))
56 mmcv.frames2video(out_path, output, fps=imgs.fps, fourcc='mp4v')
57 print('已成功导出视频 至 {}'.format(output))
58 out_dir.cleanup()