基于深度学习的篮球比赛战术数据自动采集及统计系统——5综合yolov5+deep sort+slow fast实现行为分类

前言:

现在我们手头上有了yolov5的权重模型,deepsort的权重模型,slowfast的权重模型。我们现在要做的是,输入这三个权重模型,能够实现对视频的行为分类。

基于yolov5_deepsort_slowfast的行为分类

参考项目:mirrors / wufan-tb / yolo_slowfast · GitCode

这是吴凡github项目,可以看他对项目解释,但实际上我们只要会用yolo_slowfast的内容即可。

下载下来长这样:

基于深度学习的篮球比赛战术数据自动采集及统计系统——5综合yolov5+deep sort+slow fast实现行为分类_第1张图片

注意: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文件要修改成自己对应的类别,如下

基于深度学习的篮球比赛战术数据自动采集及统计系统——5综合yolov5+deep sort+slow fast实现行为分类_第2张图片

其他函数讲解:

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! ')

基于深度学习的篮球比赛战术数据自动采集及统计系统——5综合yolov5+deep sort+slow fast实现行为分类_第3张图片

你可能感兴趣的:(深度学习,人工智能,算法)