简单地使用yolov5s训练自己的轻量级目标检测模型-基于pytorch实现

1，首先下载模型代码，点击下载

2，搭建环境

按照requirements.txt 搭建环境，这里不赘述。

3，修改dataset_directory文件夹中的内容，下图为目录结构：

为了方便展示目录树，各个文件夹我只保留了3条数据，annotations文件夹是标注好的xml文件，每个文件内容类似为：

config_file文件夹下的需修改data.yaml文件，model.yaml不用修改。data.yaml文件内容为：

names修改为你目标检测任务需要的分类的名称，nc为一共的类别数。改完即可。

images文件夹是存放原图像的文件夹。

lables文件夹是存放标签的文件夹，每个文件下的内容类似这样：

txt_file文件夹存放了训练集图像路径train.txt和验证集图像路径val.txt，二个txt文件下的内容均类似这样：

labels.cache先不用管。

如果还没有标注数据的，推荐使用labelimg标注工具，点击下载labelimg工具。

3，修改training.py文件进行模型训练

将object_names改为自己的类别名称列表。然后运行training.py程序就可以训练模型了。

运行程序后，会自动生成runs文件夹，目录结构如下：

exp10_dataset_directory为最后保存的模型。

4，修改detect.py文件进行推理

detect.py文件内容为:

根据自己的需要，修改detect_image函数中inputfile，以及out_to_file，注意修改path_to_weight为自己训练好的模型，路径为：runs/exp10_dataset_directory/weights/best_dataset_directory.pt。

注意：如果不修改out_to_file，默认保存结果为路径为：images/detected文件夹下。

修改完之后就可以运行程序进行推理了，如果要检测视频，请修改detect_video函数。

Done!!!