【人工智能学习之STGCN训练自己的数据集】

STGCN训练自己的数据集

准备事项

1. st-gcn代码下载与环境配置

git clone https://github.com/yysijie/st-gcn.git
cd st-gcn
pip install -r requirements.txt
cd torchlight
python setup.py install
cd ..

1. 数据集结构 可以使用open pose制作数据集，制作过程见下章。 （参考openpose环境搭建和利用openpose提取自建数据集）

我的数据集结构如下：

dataset/
├── stgcn_data/# 最终使用的数据集
│   ├── train                      # 训练数据集
│   └── val/# 验证数据集
├── video/# 视频数据集
│   └── fall                      # 分类0视频文件夹
│    └── normal                    # 分类1视频文件夹
│    └── resized                    # 视频缩放与json
│        └── data                # 视频对应的json
│        └── fall                # 分类0缩放视频
│        └── normal                # 分类1缩放视频
│        └── snippets            # 视频每一帧的json
│    └── label0.txt                # 分类0标签文本
│    └── label1.txt                # 分类1标签文本
└── pose_demo/# openpose目录
    └── bin/# 目录
        └── openpose_demo.exe/# openpose的exe可执行文件
        └── XXX.dll/# 各种依赖项

数据集制作

视频转json

首先需要将视频放到对应的目录下，目录名称就是你的类名。

提示一下：用于st-gcn训练的数据集视频帧数不要超过300帧，5~6s的视频时长比较好，不要10几s的视频时长。要不然会报：index 300 is out of bounds for axis 1 with size 300 这种错误。很多博主使用的用FFmpeg对视频进行缩放，但我FFmpeg老是出问题，索性直接自己用cv2实现了。
【win64中FFmpegReader报错：Cannot find installation of real FFmpeg (which comes with ffprobe).】

以下是我video2json的代码，存于dataset目录下运行：

#!/usr/bin/env pythonimport os
import argparse
import json
import shutil

import numpy as np
import torch
# import skvideo# ffmpeg_path = r"/dataset/ffmpeg-master-latest-win64-gpl/ffmpeg-master-latest-win64-gpl/bin"# skvideo.setFFmpegPath(ffmpeg_path)import skvideo.io

# from .io import IOimport tools
import tools.utils as utils

import cv2
import os

defresize_video(input_path, output_path, size=(340,256), fps=30):# 打开视频文件
    cap = cv2.VideoCapture(input_path)# 获取视频帧率if fps ==0:
        fps = cap.get(cv2.CAP_PROP_FPS)# 获取视频编码格式
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')# 创建 VideoWriter 对象
    out = cv2.VideoWriter(output_path, fourcc, fps, size)whileTrue:
        ret, frame = cap.read()ifnot ret:break# 调整帧大小
        resized_frame = cv2.resize(frame, size)# 写入输出文件
        out.write(resized_frame)# 释放资源
    cap.release()
    out.release()defresize_all_video(originvideo_file,resizedvideo_file):# 获取视频文件名列表
    videos_file_names =[f for f in os.listdir(originvideo_file)if f.endswith('.mp4')]# 遍历并处理每个视频文件for file_name in videos_file_names:
        video_path = os.path.join(originvideo_file, file_name)
        outvideo_path = os.path.join(resizedvideo_file, file_name)
        resize_video(video_path, outvideo_path)print(f'{file_name} resize success')defget_video_frames(video_path):"""
    读取视频文件并返回所有帧
    :param video_path: 视频文件路径
    :return: 帧列表
    """
    frames =[]
    cap = cv2.VideoCapture(video_path)whileTrue:
        ret, frame = cap.read()ifnot ret:break
        frames.append(frame)

    cap.release()return frames

classPreProcess():"""
        利用openpose提取自建数据集的骨骼点数据
    """defstart(self):###########################修改处################
        type_number =2
        gongfu_filename_list =['fall','normal']#################################################for process_index inrange(type_number):

            gongfu_filename = gongfu_filename_list[process_index]# 标签信息# labelgongfu_name = 'xxx_{}'.format(process_index)
            labelAction_name ='{}'.format(gongfu_filename)
            label_no = process_index

            # 视频所在文件夹
            originvideo_file ='./video/{}/'.format(gongfu_filename)# resized视频输出文件夹
            resizedvideo_file ='./video/resized/{}/'.format(gongfu_filename)'''
            videos_file_names = os.listdir(originvideo_file)

            # 1. Resize文件夹下的视频到340x256 30fps
            for file_name in videos_file_names:
                video_path = '{}{}'.format(originvideo_file, file_name)
                outvideo_path = '{}{}'.format(resizedvideo_file, file_name)
                writer = skvideo.io.FFmpegWriter(outvideo_path,
                                                 outputdict={'-f': 'mp4', '-vcodec': 'libx264', '-s': '340x256',
                                                             '-r': '30'})
                reader = skvideo.io.FFmpegReader(video_path)
                for frame in reader.nextFrame():
                    writer.writeFrame(frame)
                writer.close()
                print('{} resize success'.format(file_name))
'''# 1. Resize文件夹下的视频到340x256 30fps cv处理
            resize_all_video(originvideo_file,resizedvideo_file)# 2. 利用openpose提取每段视频骨骼点数据
            resizedvideos_file_names = os.listdir(resizedvideo_file)for file_name in resizedvideos_file_names:
                outvideo_path ='{}{}'.format(resizedvideo_file, file_name)# openpose = '{}/examples/openpose/openpose.bin'.format(self.arg.openpose)
                openpose ='C:/WorkFiles/company_server_SSH/st-gcn-master/dataset/pose_demo/bin/OpenPoseDemo.exe'
                video_name = file_name.split('.')[0]
                output_snippets_dir ='./video/resized/snippets/{}'.format(video_name)print(f"Output snippets directory: {output_snippets_dir}")
                output_sequence_dir ='video/resized/data'
                output_sequence_path ='{}/{}.json'.format(output_sequence_dir, video_name)

                label_name_path ='video/label{}.txt'.format(process_index)withopen(label_name_path)as f:
                    label_name = f.readlines()
                    label_name =[line.rstrip()for line in label_name]# pose estimation
                openpose_args =dict(
                    video=outvideo_path,
                    write_json=output_snippets_dir,
                    display=0,
                    render_pose=0,
                    model_pose='COCO')
                command_line = openpose +' '
                command_line +=' '.join(['--{} {}'.format(k, v)for k, v in openpose_args.items()])print(f"Running command: {command_line}")
                shutil.rmtree(output_snippets_dir, ignore_errors=True)
                os.makedirs(output_snippets_dir)
                os.system(command_line)# pack openpose ouputs# video = utils.video.get_video_frames(outvideo_path)
                video = get_video_frames(outvideo_path)
                height, width, _ = video[0].shape

                # 这里可以修改label, label_index
                video_info = utils.openpose.json_pack(
                    output_snippets_dir, video_name, width, height, labelAction_name , label_no)ifnot os.path.exists(output_sequence_dir):
                    os.makedirs(output_sequence_dir)withopen(output_sequence_path,'w')as outfile:
                    json.dump(video_info, outfile)iflen(video_info['data'])==0:print('{} Can not find pose estimation results.'.format(file_name))returnelse:print('{} pose estimation complete.'.format(file_name))if __name__ =='__main__':
    p=PreProcess()
    p.start()

注意：
如果你的openpose装的是cpu版本，会比较慢，跑很多视频生成json文件的时间就会很长。这个时候晚上如果去睡觉了一定一定一定要设置为从不熄屏休眠！！！从不！！！
不然就会这样：

经过漫长的等待之后。。。。终于运行完成了，我们可以得到

1. 缩放过后的视频。
2. 每一帧的json：在snippets目录下，以视频名称命名的文件下的每一帧json文件：打开是这样的：
3. 每一个视频的json：在data目录下，以视频名称命名的json文件：打开是这样的： 到这里，得到了每个视频的json文件，数据集的制作就完成了一半了。

jsons转json

得到了每个视频的json文件之后，需要我们手动将data目录下的json文件分配到train和val下，一般按照9：1划分。放好之后运行下面这段jsons2json.py:

以下是我jsons2json的代码，存于stgcn_data目录下运行：

import json
import os

if __name__ =='__main__':
    train_json_path ='./train'
    val_json_path ='./val'
    test_json_path ='./test'

    output_train_json_path ='./train_label.json'
    output_val_json_path ='./val_label.json'
    output_test_json_path ='./test_label.json'

    train_json_names = os.listdir(train_json_path)
    val_json_names = os.listdir(val_json_path)
    test_json_names = os.listdir(test_json_path)

    train_label_json =dict()
    val_label_json =dict()
    test_label_json =dict()for file_name in train_json_names:
        name = file_name.split('.')[0]
        json_file_path ='{}/{}'.format(train_json_path, file_name)
        json_file = json.load(open(json_file_path))

        file_label =dict()iflen(json_file['data'])==0:
            file_label['has_skeleton']=Falseelse:
            file_label['has_skeleton']=True
        file_label['label']= json_file['label']
        file_label['label_index']= json_file['label_index']

        train_label_json['{}'.format(name)]= file_label

        print('{} success'.format(file_name))withopen(output_train_json_path,'w')as outfile:
        json.dump(train_label_json, outfile)for file_name in val_json_names:
        name = file_name.split('.')[0]
        json_file_path ='{}/{}'.format(val_json_path, file_name)
        json_file = json.load(open(json_file_path))

        file_label =dict()iflen(json_file['data'])==0:
            file_label['has_skeleton']=Falseelse:
            file_label['has_skeleton']=True
        file_label['label']= json_file['label']
        file_label['label_index']= json_file['label_index']

        val_label_json['{}'.format(name)]= file_label

        print('{} success'.format(file_name))withopen(output_val_json_path,'w')as outfile:
        json.dump(val_label_json, outfile)for file_name in test_json_names:
        name = file_name.split('.')[0]
        json_file_path ='{}/{}'.format(test_json_path, file_name)
        json_file = json.load(open(json_file_path))

        file_label =dict()iflen(json_file['data'])==0:
            file_label['has_skeleton']=Falseelse:
            file_label['has_skeleton']=True
        file_label['label']= json_file['label']
        file_label['label_index']= json_file['label_index']

        test_label_json['{}'.format(name)]= file_label

        print('{} success'.format(file_name))withopen(output_test_json_path,'w')as outfile:
            json.dump(test_label_json, outfile)

运行完成后，我们就可以得到总的训练和验证的json文件了，里面包含了所有的视频动作关键点和标签信息。
还差一步就完成啦！

json转npy&pkl

第三步需要用到官方的工具文件kinetics_gendata.py。
所在的路径应该是：./st-gcn-master/tools/kinetics_gendata.py

这里有两处地方需要修改：

1. 修改考虑参数

defgendata(
        data_path,
        label_path,
        data_out_path,
        label_out_path,
        num_person_in=3,#observe the first 5 persons#这个参数指定了在每个视频帧中考虑的最大人数。#例如，如果设置为5，则脚本会尝试从每个视频帧中获取最多5个人的骨架信息。#这并不意味着每个视频帧都会有5个人，而是说脚本最多会处理5个人的数据。
        num_person_out=1,#then choose 2 persons with the highest score#这个参数指定了最终每个序列（或视频）中保留的人数。#即使输入中有更多的人员，也只有得分最高的两个人会被选择出来用于训练。#这里的“得分”通常指的是骨架检测的置信度分数，即算法对某个点确实是人体某部位的信心程度。
        max_frame=300):

1. 修改路径位置

if __name__ =='__main__':
    parser = argparse.ArgumentParser(
        description='Kinetics-skeleton Data Converter.')
    parser.add_argument('--data_path', default='./st-gcn-master/dataset/stgcn_data')
    parser.add_argument('--out_folder', default='./st-gcn-master/dataset/stgcn_data')
    arg = parser.parse_args()

    part =['train','val']for p in part:
        data_path ='{}/{}'.format(arg.data_path, p)
        label_path ='{}/{}_label.json'.format(arg.data_path, p)
        data_out_path ='{}/{}_data.npy'.format(arg.out_folder, p)
        label_out_path ='{}/{}_label.pkl'.format(arg.out_folder, p)ifnot os.path.exists(arg.out_folder):
            os.makedirs(arg.out_folder)
        gendata(data_path, label_path, data_out_path, label_out_path)

ok，检查一下是否得到了npy和pkl文件

yes，到这里就可以开始训练啦。

训练STGCN

添加图结构

在net/utils/graph.py文件里面get_edge函数中保存的是不同的图结构。

注意这里的默认的layout如果符合自己定义的姿态就不用修改，否则需要自定义一个，本文采用的openpose即默认的openpose的18个关键点，不需要修改。其中num_node为关键点的个数，neighbor_link为关键点连接关系。如果自己的数据集是新定义的姿态点数不为18，在后续转换中可能还有修改需要保持一致。

修改训练参数

1. 将config/st_gcn/kinetics-skeleton/train.yaml复制一份到根目录，重命名为mytrain.yaml，并修改其中参数。
2. data_path和label_path修改为之前生成的文件路径；
3. num_class改为自建数据集的行为类别个数；
4. layout参数修改为之前添加的layout类别；
5. strategy设置为spatial；
6. 修改使用的GPU数量，单个设置device: [0]；
7. optim部分适当调整，base_lr: 0.1是基础学习率，step: [80, 120, 160, 200]：这表明使用了一种学习率衰减策略，num_epoch: 200：指定了整个训练过程将进行多少个周期（epochs）
8. 不知道是不是我搞错了，居然不会自动保存best？而是每10轮自动保存一次；

以下是我的yaml：

work_dir:./work_dir/recognition/kinetics_skeleton/ST_GCN

# feeder
feeder: feeder.feeder.Feeder
train_feeder_args:
  random_choose:True
  random_move:True
  window_size:150 
  data_path: C:/WorkFiles/company_server_SSH/st-gcn-master/dataset/stgcn_data/train_data.npy
  label_path: C:/WorkFiles/company_server_SSH/st-gcn-master/dataset/stgcn_data/train_label.pkl
test_feeder_args:
  data_path: C:/WorkFiles/company_server_SSH/st-gcn-master/dataset/stgcn_data/val_data.npy
  label_path: C:/WorkFiles/company_server_SSH/st-gcn-master/dataset/stgcn_data/val_label.pkl

# model
model: net.st_gcn.Model
model_args:
  in_channels:3
  num_class:2
  edge_importance_weighting:True
  graph_args:
    layout:'openpose'
    strategy:'spatial'# training
device:[0]
batch_size:32
test_batch_size:32#optim
base_lr:0.1
step:[80,120,160,200]
num_epoch:200

开始训练

训练指令（记得先激活环境，或者在pychar的终端运行）
yaml文件可以直接放在根目录下；也可以放在自己喜欢的位置（指令需要加上路径）

python main.py recognition -c mytrain.yaml

这里我的验证集俩分类各取了3个，总训练集200多个，所以top很高。

测试

可以像训练集那样仿照一个测试集出来。
也可以像我一样直接使用视频进行测试，不过我的代码涉及了一些实际的工业工作内容，无法提供，但可以给出一些思路：

1. 通过其他关键点网络输出关键点(我用的YoloPose)
2. 将某自定义时长内的所有关键点拼接在一起重塑为 (N, in_channels, T_in, V_in, M_in) 形状
3. 通过网络输出获取这个时间段内的行为分类