现在我们手头上有了yolov5的权重模型,deepsort的权重模型,slowfast的权重模型。我们现在要做的是,输入这三个权重模型,能够实现对视频的行为分类。
参考项目:mirrors / wufan-tb / yolo_slowfast · GitCode
这是吴凡github项目,可以看他对项目解释,但实际上我们只要会用yolo_slowfast的内容即可。
下载下来长这样:
注意:1.可以自行导入yolov5文件夹,2.新建slowfast_weight文件夹,存放slowfast权重文件
(可参考文章:Yolov5+Deepsort+Slowfast实现实时动作检测_Lukarov的博客-CSDN博客)
# 加载yolov5权重模型
model = torch.hub.load('D:\deep_sort\yolov5', 'custom',
path_or_model=r'D:\yolo_slowfast\yolov5\runs\train\exp7\weights\best.pt', source='local')
cudnn.benchmark = True ##设置为True
model.conf = config.conf
model.iou = config.iou
model.max_det = 200
if config.classes:
model.classes = config.classes
device = config.device
imsize = config.imsize
# 加载官方ava2.2权重模型
# video_model = slowfast_r50_detection(True).eval().to(device)
# 加载自己训练的slowfast权重模型
video_model = slowfast_r50_detection(False).eval().to(device)
video_model.load_state_dict(
torch.load(r'D:\yolo_slowfast\slowfast_weight\checkpoint_epoch_00020.pyth')['model_state'], strict=False)
# 加载deepsort权重模型
deepsort_tracker = DeepSort("deep_sort_pytorch/deep_sort/deep/checkpoint/ckpt.t7")
# 加载动作名称和id信息
ava_labelnames, _ = AvaLabeledVideoFramePaths.read_label_map("selfutils/temp.pbtxt")
注意:selfutils/temp.pbtxt中的动作+id文件要修改成自己对应的类别,如下
def save_yolopreds_tovideo(yolo_preds, id_to_ava_labels, color_map, output_video, ID_target, ID_target_cnt):
for i, (im, pred) in enumerate(zip(yolo_preds.imgs, yolo_preds.pred)):
im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
if pred.shape[0]:
for j, (*box, cls, trackid, vx, vy) in enumerate(pred):
if int(cls) != 0:
ava_label = ''
elif trackid in id_to_ava_labels.keys():
ava_label = id_to_ava_labels[trackid].split(' ')[0]
else:
# ava_label = 'Unknow'
ava_label = 'positioning'
# trackid 记录跟踪序列号 int
# yolo_preds.names[int(cls)] 记录球员id str
# ave_label 记录动作类型 str
text = '{} {} {}'.format(int(trackid), yolo_preds.names[int(cls)], ava_label)
color = color_map[int(cls)]
im = plot_one_box(box, im, color, text)
####在这里统计动作类型指标####
if ID_target[yolo_preds.names[int(cls)]] != ava_label:
ID_target[yolo_preds.names[int(cls)]] = ava_label # 更新当前的动作
ID_target_cnt[yolo_preds.names[int(cls)]][ava_label] += 1 # 当前的动作指标数目加一
output_video.write(im.astype(np.uint8))
这段函数中,可以获得行为分类识别出的结果信息,用于统计指标
视频结果:形成的视频结果可以自行修改路径,如上’--output'。输入预测视频即'--input'
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--input', type=str, default=r"D:\basketballvideo\MyVideo_1.mp4",
help='test imgs folder or video or camera')
parser.add_argument('--output', type=str, default="D:\yolo_slowfast\Output\output2.mp4",
help='folder to save result imgs, can not use input folder')
# object detect config
parser.add_argument('--imsize', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf', type=float, default=0.1, help='object confidence threshold')
parser.add_argument('--iou', type=float, default=0.1, help='IOU threshold for NMS')
parser.add_argument('--device', default=0, help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--classes', default=0, nargs='+', type=int,
help='filter by class: --class 0, or --class 0 2 3')
config = parser.parse_args()
print(config)
main(config)
由于只训练了40张数据,所以追踪效果不是太好,可以看到全程只追踪到了球员1的跑位动作
战术数据结果: 调用out_target_csv函数(自己编写的)
#在这里导出指标文件
path = 'D:\yolo_slowfast\out.csv'
out_target_csv(path,ID_target_cnt)
outputvideo.release()
print('saved target.csv to:', path)
print('saved video to:', vide_save_path)
print(' Done! ')