这两天特别认真的分析了一下在该比赛中使用YOLO模型进行训练和预测的代码逻辑,以及一些具体代码的意义。
整个代码主要分为几个部分:
-
1. 导入各类库、下载数据的训练集和测试集,并进行文件的解压缩
-
2.数据读取:主要拿一个视频第一帧为例子,为我们详细分析了每一个json文件和视频每一帧的基本情况
-
3. 数据转换:
-
创建dataset,train,val的文件夹,以便后期模型训练时存储信息
-
创建一个yolo.yaml文件,写入数据集的路径和训练、验证的路径,以及检测类别,以便后期模型进行训练时,可以直接从该信息中读取基本信息,并找到数据集的路径
-
找到所有json和MP4文件,进行排序,并设置检测类别
-
取出前5个视频和对应的json文件
-
首先,读取每一个视频文件里的每一帧,并将其按照顺序保存成jpg格式
-
其次,调出与每一帧索引相对应的json文件里面的frame_id,如果这个文件有信息,则说明该帧里面有违规类别;于是,我们需要创建一个相对应的.txt文件,写入该帧中的几个参数:违规类别,x_center,y_center,边界框的宽度,边界框的高度(记得要进行归一化)
-
4. 模型训练:
所有训练集以及配置文件准备好后,终于可以进行模型的训练了,具体包含以下几步:
- 下载我们需要使用的YOLO模型,创建所需目录和文件
-
模型开始训练啦!其中数据集路径可以从yolo.yaml文件中读取
-
5. 使用训练模型进行违规类型的预测
-
创建一个result目录,以便存储我们最终的预测结果
-
导入我们在之前模型训练中表现最好的权重参数文件,以进行测试集的预测
以下整个循环的逻辑:
-
首先针对每个视频文件处理,设置空的submit_json列表,并使用训练好的模型进行预测
-
然后在每个视频文件内部,针对视频的每一帧进行目标检测和信息提取操作
-
将所有信息添加进submit_json列表中
-
最后,创建一个json文件,将submit_json列表内容写入up,并转换为Json文件格式
-
最终,删除与result.zip相关的文件,防止有其他文件和我们最终的结果文件重名
-
将result/目录下的文件压缩,以result.zip命名