如何制作YOLOv3模型训练数据集？

上一篇笔记YOLOv3网络结构和基础知识学习笔记我们介绍了YOLO V3的网络结构,yolov3的输入输出，yolo坐标轴如何进行反算。这里，我们一起讨论一下如何使用YOLO V3网络。

我的其他笔记链接：

1. 使用K-means算法寻找yolo的锚框 - 简书 (jianshu.com)
2. 如何将图片和锚框一起resize——python代码实现 - 简书 (jianshu.com)
   3.YOLO学习笔记——目标检测算法的简介（RCNN, SSD, YOLO) - 简书 (jianshu.com)
   4.YOLOv3网络结构和基础知识学习笔记 - 简书 (jianshu.com)
3. 如何制作YOLOv3模型训练数据集？ - 简书 (jianshu.com)
4. 如何训练YOLOv3模型？pytorch代码实现 - 简书 (jianshu.com)
5. YOLOv3、YOLOv4、YOLOv5之间的区别 - 简书 (jianshu.com)

一、制作数据集

在训练模型之前，我们需要制作数据集和标签。标签的形状应该和神经网络的输出应该一致，这样才能计算损失。那么我们怎样将标签从原来的一维 [图片路径，分类，x1, y1, x2, y2 ...] 制作成 (N, 45, 13, 13) 呢？

首先，如果我们坐标是左上角右小角的形式，可以先转换成中心点和宽高的方式。

一个小格子代表的是原图中的一个小窗口，然后我们的标签信息要表示在通道上，然后我们可以将分类转换为one-hot标签，然后就再将置信度和坐标位置转换成偏移量，最后将所有的数据填写到对应的通道上。

那么，中心点和onehot编码都是好理解的，那么置信度怎么填，填到上面哪个特征图上呢？因为三个特征图分别代表三种不同形状的锚框。我们肯定是希望填在那个和真实框相似度高的特征图上。什么可以代表两个框的置信度呢？答案是：IOU值！那么我们就可以让真实框分别和三个锚框求IOU，用求的IOU值作为置信度，填写到每一个特征图上。三个特征图上，one-hot编码和中心点的偏移量是一样的，其他的都不一样。

可能有同学说，为什么三个特征图都要填写，而不是不只填写IOU最高的那个呢？其他特征图的置信度填写为0就好啦？是因为，在特征图上，那个窗口里确实存在目标，你不能说它不存在吧，只是说锚框不合适导致了置信度降低嘛。

上面是基于13*13的特征图描述的，那么我们还要制作出26*26、52*52的特征图。

我们现在以这个百度AI开放的数据集为例。
车辆识别 - 飞桨AI Studio (baidu.com)

（1）第一步将图片都缩放成416*416大小

首先我们从VOC格式的XML文件里面读取相应信息，将其保存为txt格式的文件，接着我们要将图片进行等比缩放，并且框也是要等比缩放的。其代码实现如下：

import glob
import xml.etree.ElementTree as ET
import numpy as np
import os 
import PIL.Image as Image
import PIL.ImageDraw as Draw
import tqdm

xml_path = r"./dataset/car-identify/car-main/dataset/dataset/annotation/"
img_path = r"./dataset/car-identify/car-main/dataset/dataset/images/"
txt_path = r"./dataset/car-identify/car-main/dataset/dataset/Imagesets/img_label.txt"
img_dst_path = r"./dataset/car-identify/car-main/dataset/dataset/anno_img/"
label = {
    "truck":1,
    "bus":2,
    "SUV":3,
    "taxi":4,
    "car":5
}

def read_XML():
    lines = []
    try:
        file_list = os.listdir(xml_path)
        for _xml in file_list:
            line = []
            xml_file = xml_path + _xml
            with open(xml_file) as f:
                tree = ET.parse(f)
                height = int(tree.findtext('./size/height'))
                width = int(tree.findtext('./size/width'))
                if height <= 0 or width <= 0:
                    continue

                file_name = tree.findtext('./filename')
                line.append(file_name)

                # 对于每一个目标都获得它的宽高
                for obj in tree.iter('object'):
                    cat = label[str(obj.findtext('name'))]
                    xmin = int(float(obj.findtext('bndbox/xmin')))
                    ymin = int(float(obj.findtext('bndbox/ymin')))
                    xmax = int(float(obj.findtext('bndbox/xmax')))
                    ymax = int(float(obj.findtext('bndbox/ymax')))

                    line.extend([cat, xmin, ymin, xmax, ymax])
            f.close()
            lines.append(line)
    except Exception as e:
        # print("XML FILE OPEN ERROR!")
        print(e)
    return lines

"""
将图像resize， 并且将他的box也resize
"""
def img_resize(xml_data, size):
    # 读取图像
    f = open(txt_path, "w")
    for line in tqdm.tqdm(xml_data):
        bg_img = Image.new('RGB',size,(0,0,0))
        # 先将图片进行缩放
        path = img_path + line[0]
        img = Image.open(path)
        iw,ih = img.size
        max_side = max(iw,ih)
        scale = max_side / max(size)
        img = img.resize((int(iw/scale), int(ih/scale)))
        if iw > ih:
            dy = int((size[0] - ih/scale)/2)
            dx = 0
            bg_img.paste(img,(0,dy))
        else:
            dx = int((size[0] - iw/scale)/2)
            dy = 0
            bg_img.paste(img,(dx,0))

        # 这里是变化后的
        strs = f"{line[0]} "
        # 将框也对应缩放,读取每一个框
        for i in range(1,len(line[1:]),5):
            box = line[i:i+5] # [1,6)、[6,11)
            box = np.array(box[1:],dtype=np.float32)
            box = box / scale
            box[0] = box[0]+dx
            box[1] = box[1]+dy
            box[2] = box[2]+dx
            box[3] = box[3]+dy

            w = int(box[2] - box[0])
            h = int(box[3] - box[1])
            cx = int(box[0] + w/2)
            cy = int(box[1] + h/2)

            # box[0] = cx - w/2
            # box[1] = cy - h/2
            # box[2] = cx + w/2
            # box[3] = cy + h/2

            strs += f"{line[i]} {cx} {cy} {w} {h}"
            draw = Draw.Draw(bg_img)
            draw.rectangle((box[0], box[1], box[2], box[3]),width=2,outline=(0,0,255))
            bg_img.show()
            exit()
        
        img_save_path = img_dst_path + line[0]
        bg_img.save(img_save_path)
        f.write(strs + "\n")

if __name__ == "__main__":
    xml_data = read_XML()
    img_resize(xml_data,(416,416))

（2）使用K-means生成9个锚框。Kmeans算法生成锚框的代码：

import glob
import random
import xml.etree.ElementTree as ET
import numpy as np
import PIL.Image as Image
import PIL.ImageDraw as D

# 计算IOU
def cas_iou(box, cluster):
    x = np.minimum(cluster[:, 0], box[0])
    y = np.minimum(cluster[:, 1], box[1])

    intersection = x * y
    area1 = box[0] * box[1]
    area2 = cluster[:, 0] * cluster[:, 1]
    iou = intersection / (area1 + area2 - intersection)
    return iou

# 计算平均IOU
def avg_iou(box, cluster):
    return np.mean([np.max(cas_iou(box[i], cluster)) for i in range(box.shape[0])])

def kmeans(box, k):
    # 取出一共有多少框
    row = box.shape[0]
    # 每个框各个点的位置
    distance = np.empty((row, k)) # [699, 9]

    # 最后的聚类位置
    last_clu = np.zeros((row,)) # [699,]

    np.random.seed()

    # 随机选9个当聚类中心
    cluster = box[np.random.choice(row, k, replace=False)] # [9,2]
    # cluster = random.sample(row, k)
    while True:
        # 计算每一行距离五个点的iou情况。
        for i in range(row):
            distance[i] = 1 - cas_iou(box[i], cluster)

        # 取出最小点的索引值
        near = np.argmin(distance, axis=1) # [699, 1]

        # 算法结束条件
        if (last_clu == near).all():
            break

        # 求每一个类的中位点，
        for j in range(k):
            cluster[j] = np.median(box[near == j], axis=0) # 计算中位数 [9,2]

        last_clu = near
    return cluster

# 从xml文件中读取
def load_data(path):
    data = []
    # 对于每一个xml都寻找box
    for xml_file in glob.glob('{}/*xml'.format(path)):
        tree = ET.parse(xml_file)
        height = int(tree.findtext('./size/height'))
        width = int(tree.findtext('./size/width'))
        if height <= 0 or width <= 0:
            continue

        # 对于每一个目标都获得它的宽高
        for obj in tree.iter('object'):
            xmin = int(float(obj.findtext('bndbox/xmin'))) / width
            ymin = int(float(obj.findtext('bndbox/ymin'))) / height
            xmax = int(float(obj.findtext('bndbox/xmax'))) / width
            ymax = int(float(obj.findtext('bndbox/ymax'))) / height

            xmin = np.float64(xmin)
            ymin = np.float64(ymin)
            xmax = np.float64(xmax)
            ymax = np.float64(ymax)
            # 得到宽高
            data.append([xmax - xmin, ymax - ymin])
    return np.array(data)

# 从txt中读取
def load_data2(path, img_path):
    res = []
    with open(path) as f:
        for line in f.readlines():
            arr = line.strip().split(" ")
            box = np.array(arr[2:], dtype=np.float64)
            img_file_path = img_path + arr[0]
            w,h = Image.open(img_file_path).size

            x1 = (box[0] - box[2]/2) / w
            y1 = (box[1] - box[3]/2) / h
            x2 = (box[0] + box[2]/2) / w
            y2 = (box[1] + box[3]/2) / h
            res.append([x2-x1, y2-y1])
    return np.array(res)

if __name__ == '__main__':
    SIZE = 416
    anchors_num = 9
    # 载入数据集，可以使用VOC的xml
    img_path = r"./dataset/car-identify/car-main/dataset/dataset/anno_img/"
    path = r"./dataset/car-identify/car-main/dataset/dataset/Imagesets/img_label.txt"
    path2 = r"./dataset/car-identify/car-main/dataset/dataset/annotation"

    # 载入所有的xml
    # 存储格式为转化为比例后的width,height
    data = load_data2(path, img_path)

    # 使用k聚类算法
    out = kmeans(data, anchors_num)
    out = out[np.argsort(out[:, 0])]
    print('acc:{:.2f}%'.format(avg_iou(data, out) * 100))
    print(out * SIZE)
    data = out * SIZE # [9, 2]

    # 这里是按照面积大小进行排序
    area = data[:,0] * data[:, 1]
    sort_index = np.argsort(area) # 获取面积由小到大的排序
    data = data[sort_index]
    new_data = []
    # 每三个为同一个尺寸：小框、中框、大框
    # 再按照宽高比进行排序，排完之后：竖框、方框、横框
    for i in range(0,9,3):
        boxes = data[i:i+3]
        ratio = boxes[:, 0] / boxes[:, 1]
        ratio_index = np.argsort(ratio)
        boxes = boxes[ratio_index]
        [new_data.append(box) for box in boxes]
    data = new_data

    f = open("./param/neuron_anchors.txt", 'w')
    row = np.shape(data)[0]
    for i in range(row):
        if i == 0:
            x_y = "%d,%d" % (data[i][0], data[i][1])
        else:
            x_y = ", %d,%d" % (data[i][0], data[i][1])
        f.write(x_y)
    f.close()

这里也直接给出我计算出来的九个建议框：

（3）构建数据集Dataset

import torch 
import numpy as np
import PIL.Image as Image
from torch.utils.data import Dataset
from torchvision.transforms import ToTensor
import cfg
import math 
from torch.nn.functional import one_hot

anchor_file_path = r"./param/neuron_anchors.txt"
xml_path = r"./dataset/car-main/dataset/dataset/annotation/"
img_path = r"./dataset/car-main/dataset/dataset/anno_img/"
txt_path = r"./dataset/car-main/dataset/dataset/Imagesets/img_label.txt"


class Car_Dataset(Dataset):
    def __init__(self) -> None:
        super().__init__()
        f = open(txt_path)
        self.dataset = f.readlines()
        f.close()
        self.trans = ToTensor()

    def __len__(self):
        return len(self.dataset) 

    def __getitem__(self, index):
        labels = {} # 创建一个空字典
        line = self.dataset[index].strip().split(" ")
        img = Image.open(img_path + line[0])
        img = self.trans(img)
        
        # 将锚框的位置坐标分出去
        _boxes = np.array(line[1:], dtype=np.float32)
        boxes = np.split(_boxes, len(_boxes) // 5) # 每五个值为一组划分：[类别，cx，cy，w，h]

        for feature, anchors in cfg.ANCHORS_CROUP_KMEANS.items():
            labels[feature] = np.zeros(shape=(feature,feature,3,5+cfg.CLASS_NUM), dtype=np.float32)

            for box in boxes:
                cls, cx, cy, w, h = box 
                cx_offset, cx_index = math.modf(cx * feature / cfg.IMG_W) # 求中心点的索引和偏移量
                cy_offset, cy_index = math.modf(cy * feature / cfg.IMG_H) # 求中心点的索引和偏移量
                for i, anchor in enumerate(anchors):
                    p_w, p_h = w / anchor[0], h / anchor[1] # 宽高偏移量
                    p_area = w * h # 真实框的面积
                    anchor_area = cfg.ANCHORS_CROUP_AERA[feature][i] # 建议框的面积
                    # 这里的iou使用的是最小面积除以最大面积
                    iou = min(p_area, anchor_area) / max(p_area, anchor_area)
                    # labels的形状（13，13，1，10）
                    # 注意是先cy后cx
                    labels[feature][int(cy_index),int(cx_index),i] = np.array([
                        iou, 
                        cx_offset, 
                        cy_offset, 
                        p_w, 
                        p_h,
                        *one_hot(torch.tensor(int(cls)), cfg.CLASS_NUM)])
        return torch.tensor(labels[13]), torch.tensor(labels[26]), torch.tensor(labels[52]), img
    
if __name__ =="__main__":
    db = Car_Dataset()
    print(db[0][0].shape)
    print(db[0][1].shape)
    print(db[0][2].shape)
    print(db[0][3].shape)