Datawhale AI 夏令营之Yolo模型

这两天特别认真的分析了一下在该比赛中使用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命名