使用Yolov5 模型，训练自己的数据集这里以目标检测BDD数据集为例，跑出检测结果
yolov5模型下载地址：GitHub - ultralytics/yolov5: YOLOv5 in PyTorch > ONNX > CoreML > TFLite
下载完需要配置环境，yolov5的环境所需的包基本都收录在requirements.txt文件中，同时在安装一下torch（因为使用的是Pytorch框架）
如果你不会配置环境，建议看一下conda命令：
查看已有的虚拟环境
conda env list 创建虚拟环境和删除虚拟环境
anaconda命令创建python版本为x.x，名字为env_name的虚拟环境。
env_name文件可以在Anaconda安装目录envs文件下找到。
# 创建
conda create -n env_name python=x.x
# 删除
conda remove -n env_name --all创建好环境后开始上面说的所需要的包开始安装到自己的环境
anaconda安装完成之后请切换到国内的源来提高下载速度 ，命令如下：使用镜像下载速度大幅度提高
conda config --remove-key channels
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.ustc.edu.cn/anaconda/pkgs/free/
conda config --add channels https://mirrors.bfsu.edu.cn/anaconda/cloud/pytorch/
conda config --set show_channel_urls yes
pip config set global.index-url https://mirrors.ustc.edu.cn/pypi/web/simple
接着就能在环境中安装yolov5所需的包，之前需要单独装torch包
pip install -r requirements.txt接下来，整个YOLOv5就可以开始使用了。
先测试一下能不能使用
python detect.py --source data/images/bus.jpg --weights pretrained/yolov5s.pt
如果是这样跑出结果了，说明你的YOLO模型已经可以使用了。
接下来就是训练自己的数据集了
这里yolov5也提供了一些指导方法：Train Custom Data - YOLOv5 Documentation
首先看自己数据集格式的标签：
BDD100K给的数据集标签为 json文件，想送入YOLO模型跑就得改一下标签格式，改成coco数据集类似的 txt标签类型。
代码如下：
import os
import json
class BDD_to_YOLOv5:
def __init__(self):
self.writepath = "BDD100K/labels/trains/"
self.bdd100k_width_ratio = 1.0/1920
self.bdd100k_height_ratio = 1.0/1080
self.select_categorys = ["person", "rider", "car", "bus", "truck", "bike","motor", "traffic light", "traffic sign", "train"]
self.categorys_nums = {
"person": 0,
"rider": 1,
"car": 2,
"bus": 3,
"truck": 4,
"bike": 5,
"motor": 6,
"tl_green": 7,
"tl_red": 8,
"tl_yellow": 9,
"tl_none": 10,
"traffic sign": 11,
"train": 12
}
def bdd_to_yolov5(self, path):
lines = ""
with open(path) as fp:
j = json.load(fp)
write = open(self.writepath + "%s.txt" % j["name"], 'w')
for fr in j["frames"]:
for objs in fr["objects"]:
if objs["category"] in self.select_categorys:
temp_category=objs["category"]
if (temp_category == "traffic light"):
color = objs["attributes"]["trafficLightColor"]
temp_category="tl_"+color
idx = self.categorys_nums[temp_category]
cx = (objs["box2d"]["x1"] + objs["box2d"]["x2"]) / 2.0
cy = (objs["box2d"]["y1"] + objs["box2d"]["y2"]) / 2.0
w = objs["box2d"]["x2"] - objs["box2d"]["x1"]
h = objs["box2d"]["y2"] - objs["box2d"]["y1"]
if w <= 0 or h <= 0:
continue
# 根据图片尺寸进行归一化
cx, cy, w, h = cx * self.bdd100k_width_ratio, cy * self.bdd100k_height_ratio, w * self.bdd100k_width_ratio, h * self.bdd100k_height_ratio
line = f"{idx} {cx:.6f} {cy:.6f} {w:.6f} {h:.6f}\n"
lines += line
if len(lines) != 0:
write.writelines(lines)
write.close()
print("%s has been dealt!" % j["name"])
if __name__ == "__main__":
bdd_labels_dir = "BDD100K/labels/train/"
fileList = os.listdir(bdd_labels_dir)
obj = BDD_to_YOLOv5()
for path in fileList:
filepath = bdd_labels_dir+path
print(path)
obj.bdd_to_yolov5(filepath)
然后就能把数据集label全都改好了，借下来分一下训练集和验证集。
接下来改一下yaml文件
# train and val data as 1) directory: path/images/, 2) file: path/images.txt, or 3) list: [path1/images/, path2/images/]
train: cv_bdd/images/train/
val: cv_bdd/images/val/
# number of classes
nc: 13
# class names
names: ["person", "rider", "car", "bus", "truck", "bike", "motor", "tl_green", "tl_red", "tl_yellow", "tl_none", "traffic sign", "train"]这里主要要改动自己的文件读取路径，和自己相应数据集的类别，（需要自己调整）
这里yolov5提供的预训练模型有五种，根据自己需求选择。
下来就是训练自己数据集
$ python train.py --img 640 --batch 16 --epochs 5 --data coco128.yaml --weights yolov5s.pt然后就能根据跑出的结果看到训练结果了
大概实现步骤就是这样！