YOLOv5多目标识别

用于个人记录，好记性不如烂笔头

一、数据整理

​ 多目标检测其实跟单个目标检测的步骤是一样的，只是多目标检测的标签是有多个的。一张图是可以标记多个目标的。下图中，我自己的图片是一张图中只有一个目标的，我要训练一个可以识别9个目标的模型。我的标签如下：

0是1_one文件夹中所有图片的标签
1是2_two文件夹中所有图片的标签
…
…
依次类推！

具体数据整理的做法如下：

• 将图片复制到images文件夹（移动过去也可以）

• 将xml转txt格式，并保存在labels文件夹中。

• 拆分数据集，将images文件夹中的图片的路径写入到train.txt和val.txt中

代码执行后效果如下图

查看转好的txt文件

可以看到标签跟我们的文件夹是相互对应的，这里每一个txt的命名，我都给它加上对应文件夹的名字当前缀，防止重名

代码如下

# -- coding: utf-8 --
# @author : Future

import os
import glob
import re
import shutil
import tqdm

# xml转txt
def xml_to_txt(xml_path):
    with open(xml_path, "r") as f:
        data = f.read()
        xmin = int(re.findall(r'<xmin>(.*?)</xmin>', data)[0])
        ymin = int(re.findall(r'<ymin>(.*?)</ymin>', data)[0])
        xmax = int(re.findall(r'<xmax>(.*?)</xmax>', data)[0])
        ymax = int(re.findall(r'<ymax>(.*?)</ymax>', data)[0])
    cx = ((xmin + xmax) / 2) / 320  # 320在这里是图片的size
    cy = ((ymin + ymax) / 2) / 320
    w = (xmax - xmin) / 320
    h = (ymax - ymin) / 320
    return cx, cy, w, h


img_floders = r"F:\yolodata\many_numbers"
floder_list = os.listdir(img_floders)
labels_name = dict(zip(floder_list, [str(i) for i in range(len(floder_list))]))

images_path = os.path.join(img_floders, "images")
labels_path = os.path.join(img_floders, "labels")
os.makedirs(images_path, exist_ok=True, mode=0o777)
os.makedirs(labels_path, exist_ok=True, mode=0o777)

for floder_name in floder_list:
    imgs = glob.glob(os.path.join(img_floders, floder_name, "Frames","*.png"))

    for img in tqdm.tqdm(imgs):
        # r"F:\yolodata\many_numbers\1_one\Frames\6.png"
        # 重命名
        label_name = labels_name[floder_name]
        img_name = img.split("\\")[-1]
        new_img_name = "{}_{}".format(floder_name, img_name)
        new_img_path = os.path.join(images_path, new_img_name)
        xml_path = img.replace("Frames", "GroundTruth").replace(".png", ".xml")
		# 复制图片到images文件夹中
        shutil.copy(img, new_img_path)
        # 执行xml转txt操作并保存到labels文件夹
        f = open(os.path.join(labels_path, new_img_name.replace(".png", ".txt")), "w")
        cx, cy, w, h = xml_to_txt(xml_path)
        f.write("{} {} {} {} {}".format(label_name, cx, cy, w, h))
        f.close()

        
# 拆分数据集，按8:2进行拆分        
def train_val():
    imgs_path = r"F:\yolodata\many_numbers"
    imgs_list = glob.glob(os.path.join(imgs_path, "images", "*.png"))
    random.shuffle(imgs_list)

    tran_txt = open(os.path.join(imgs_path, "train.txt"), "w")
    val_txt = open(os.path.join(imgs_path, "val.txt"), "w")
    # 写入文件
    for img in imgs_list[:int(len(imgs_list) * 0.8)]:
        tran_txt.write(img + "\n")

    for img in imgs_list[int(len(imgs_list) * 0.8):]:
        val_txt.write(img + "\n")

    tran_txt.close()
    val_txt.close()

二、模型训练

​ 在yolov5工程文件中的data文件夹中创建一个yaml文件，将整理好的数据集的主路径、train.txt、val、标签和对应的名字写上

​ 修改train.py文件配置项

​ 当训练完后会提示所保存的位置，已经各种损失值以及准确率,可以看到图中提示结果保存在yolov5工程文件中的runs\train\exp

​ 打开该文件夹里面的weights，里面保存着pt模型,可以看到效果还行，我只训练了15轮

三、模型推理

（1）执行export.py，将pt模型转成onnx。

​ 做法：将数据集的yaml路径写上，并将刚刚训练好的模型best.pt（在runs\train\weight中）如果’–include’配置项没有写onnx要自己加上。运行的结果也是保存在best.pt同一个文件夹中

（2）执行推理操作

​ 下图是转好的onnx,从图中可以知道它的输入要求和输出的shape，以及标签所对应识别到的目标，所以我们只需导入onnx并对输入的图像进行处理（跟yolov5的图片处理方式要一样）即可。

输出结果中14代表着cx cy w h score label]，这里label一共有9个目标所以合起来就一共是14个

6300是它检测框的数量，因为yolov5使用三个尺寸的下采样比例分别是8,16,32，每个尺度对应的特征图大小就是

320/8=40，shape为[40x40] 320/16=20，shape为[20x20] 320/32=10，shape为[10x10]

所以一共是3x(40x40+20x20+10x10)=6300

推理代码如下

import numpy as np
import cv2
import onnxruntime as ort
import glob

provider = ort.get_available_providers()[1 if ort.get_device() == "GPU" else 0]
# 创建推理实例
ort_session = ort.InferenceSession(
    r"E:\yolomaster\yolov5\runs\train\exp\weights\best.onnx",
    providers=[provider]
)

imgs_path =r"C:\Users\ASUS\Desktop\my_test\aaaa"

for img in glob.glob(imgs_path + "/*.png"):
    img_src = cv2.imread(img)
    img_copy=img_src.copy()
    # 增加维度并归一化
    new_frame = img_copy[np.newaxis, :, :, :].astype(np.float32)
    new_frame = new_frame.transpose(0, 3, 1, 2) / 255.0
    # 推理
    result = ort_session.run(output_names=["output"],
                             input_feed={"images": new_frame})

    # 输出结果是列表嵌套一个数组，shape是 (1, 6300, 14)
    print(result[0].shape) 

    best_index = result[0][:, :, 4].argmax()
    best_reslut = result[0][:, best_index, :]
    cx, cy, w, h, score, label = best_reslut[0][0], best_reslut[0][1], best_reslut[0][2], best_reslut[0][3], \
                                 best_reslut[0][4], best_reslut[0][5]
    cv2.putText(img_copy, str(score), (int(cx - w / 2), int(cy - h / 2) - 5), cv2.FONT_HERSHEY_SIMPLEX, 1,
                (255, 255, 255), 2)
    # 绘制目标的边框
    cv2.rectangle(img_copy, (int(cx - w / 2), int(cy - h / 2)), (int(cx + w / 2), int(cy + h / 2)), (0, 0, 255),
              2)
    cv2.imshow("img",np.hstack([img_src, img_copy]))
    cv2.waitKey()
    cv2.destroyAllWindows()
    if cv2.waitKey(1) & 0xFF == 27:
        cv2.destroyAllWindows()

可以看到效果还不错

完结！！！

723872327419)]

[外链图片转存中…(img-lBls5prp-1723872327420)]