YOLOv5的使用与训练

虚拟环境终端下载YOLOv5并下载相关依赖

git clone https://github.com/ultralytics/yolov5 # clone
cd yolov5
pip install -r requirements.txt # install

使用

使用方法1

可以在yolov5下自己创建一个xxx.py文件：

import torch

# Model
model = torch.hub.load("ultralytics/yolov5", "yolov5s")  # or yolov5n - yolov5x6, custom
# Images
img = "https://ultralytics.com/images/zidane.jpg"  # or file, Path, PIL, OpenCV, numpy, list
# Inference
results = model(img)
# Results
results.print()  # or .show(), .save(), .crop(), .pandas(), etc.

使用python xxx.py运行该文件，可以使用现成的训练好的yolov5s模型进行目标检测。

使用方法2

也可以在终端输入python detect.py --weights yolov5s.pt --source xxxxxx进行检测。

例如python detect.py --weights yolov5s.pt --source 0可以实现摄像头实时目标检测。

或例如：python detect.py --weights yolov5s.pt --source img.jpg

（img.jpg对应的是位于yolov5文件夹下的名称为img.jpg的图像文件）

根据输出提示可以去yolov5/runs/detect/…下查看检测结果：

部署

将YOLOv5部署到web端（flask+js），此处省略代码，全部放在项目里了。

部署项目的app.py代码里有个地方要注意修改一下：

@app.route('/sh', methods=['GET', 'POST'])  #定义新路由，显示图片
def hello_world():
    #图片路径，推理完之后，默认保存的就是runs\\detect\\exp，这里加上filename，是变成完整的图片路径，然后才能获取显示
    img_path = 'runs/detect/exp/' + str(filename)
    img_stream = return_img_stream(img_path) #获取图片流
    return render_template('index.html', img_stream=img_stream)

这里img_path有可能要改，比如window系统下可能是：

img_path = 'runs\\detect\\exp\\' + str(filename)

但macOS是：

img_path = 'runs/detect/exp/' + str(filename)

训练自己的模型（口罩数据集）CPU版

创建一个classes.txt文件（位置根据自己的习惯，记住路径就行）：

再创建一个xml_to_txt.py文件：

import os
import xml.etree.ElementTree as ET

def convert_voc_to_yolo(xml_file, txt_file, classes):
    tree = ET.parse(xml_file)
    root = tree.getroot()

    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)

    with open(txt_file, 'w') as f:
        for obj in root.iter('object'):
            cls = obj.find('name').text
            if cls not in classes:
                continue
            cls_id = classes.index(cls)

            xmlbox = obj.find('bndbox')
            b = (int(xmlbox.find('xmin').text), int(xmlbox.find('xmax').text),
                 int(xmlbox.find('ymin').text), int(xmlbox.find('ymax').text))
            x_center = (b[0] + b[1]) / 2.0 / width
            y_center = (b[2] + b[3]) / 2.0 / height
            w = (b[1] - b[0]) / float(width)
            h = (b[3] - b[2]) / float(height)
            f.write(f"{cls_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}\n")

def convert_dataset(xml_dir, txt_dir, classes_file):
    with open(classes_file, 'r') as f:
        classes = [line.strip() for line in f.readlines()]

    if not os.path.exists(txt_dir):
        os.makedirs(txt_dir)

    for xml_file in os.listdir(xml_dir):
        if xml_file.endswith('.xml'):
            txt_file = os.path.join(txt_dir, os.path.splitext(xml_file)[0] + '.txt')
            convert_voc_to_yolo(os.path.join(xml_dir, xml_file), txt_file, classes)

# 示例用法
xml_dir = '/Users/wuleihuan/Desktop/口罩数据集/annotations'  # XML 文件目录
txt_dir = '/Users/wuleihuan/Desktop/口罩数据集/label'  # 转换后的 TXT 文件目录
classes_file = '/Users/wuleihuan/yolov5/classes.txt'  # 类别文件路径

convert_dataset(xml_dir, txt_dir, classes_file)

运行这个文件⬆️，然后在转换后的 TXT 文件目录里找到转换好的txt文件。

在yolov5/data文件夹下新建如下文件夹：

自己划分一下数据集，我把maksssksksss0-maksssksksss600当训练集，其他的当验证集了。

把训练集的png文件复制到yolov5/data/mydata/images/train文件夹里，txt文件（刚刚转换得到的）放到yolov5/data/mydata/labels/train文件里，同理，验证集的png文件放yolov5/data/mydata/images/val里，txt文件放yolov5/data/mydata/labels/val里：

（cache文件是后面过程中出现的，不用管）

在yolov5/data文件夹下新建mydata.yaml（path改成自己的）：

path: /Users/wuleihuan/yolov5/data/mydata
train: images/train
val: images/val
test:  # test images (optional)

# Classes
nc: 2  # number of classes
names: ['with_mask','without_mask']  # class names