如何制作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)
最后编辑于
©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 214,837评论 6 496
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 91,551评论 3 389
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 160,417评论 0 350
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 57,448评论 1 288
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 66,524评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 50,554评论 1 293
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,569评论 3 414
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 38,316评论 0 270
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,766评论 1 307
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 37,077评论 2 330
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 39,240评论 1 343
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,912评论 5 338
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,560评论 3 322
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 31,176评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,425评论 1 268
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 47,114评论 2 366
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 44,114评论 2 352

推荐阅读更多精彩内容