基于YOLOv5的钢材表面缺陷检测

配置环境

conda create -n yolo
pip install -r requirements.txt -i https://repo.huaweicloud.com/repository/pypi/simple # 亲测华为云最稳定

# 如果在下载时提示timeout，试试这句命令：
pip --default-timeout=100 install -r requirements.txt -i https://repo.huaweicloud.com/repository/pypi/simple

1、将xml文件转为符合yolo的txt文件，代码如下：

import xml.etree.ElementTree as ET
import os


# box [xmin,ymin,xmax,ymax]
def convert(size, box):
    x_center = (box[2] + box[0]) / 2.0
    y_center = (box[3] + box[1]) / 2.0
    # 归一化
    x = x_center / size[0]
    y = y_center / size[1]
    # 求宽高并归一化
    w = (box[2] - box[0]) / size[0]
    h = (box[3] - box[1]) / size[1]
    return (x, y, w, h)


# xml2yolo
def convert_annotation(xml_paths, yolo_paths, classes):
    xml_files = os.listdir(xml_paths)
    # 生成无序文件列表
    for file in xml_files:

        xml_file_path = os.path.join(xml_paths, file)
        yolo_txt_path = os.path.join(yolo_paths, file.split(".")[0]
                                     + ".txt")
        tree = ET.parse(xml_file_path)
        root = tree.getroot()
        size = root.find("size")
        # 获取xml的width和height的值
        w = int(size.find("width").text)
        h = int(size.