自己写—YOLOv3(2)—加载数据集

808 阅读9分钟

携手创作,共同成长!这是我参与「掘金日新计划 · 8 月更文挑战」的第9天,点击查看活动详情

前言

这个类主要目的是将用于训练数据处理为模型可以接受数据格式,然后在训练过程中,不断按事先指定批量大小将数据喂给模型,模型用这一个批量数据进行训练,也就是更新一次参数。这算一次迭代。这里我们定义数据集类是继承于 PyTorch 提供的 Dataset 这个基类,然后需要通过定义 __len____getitem__ 这两个协议接口来实现将数据集处理方式,关于对于图像数据和标注处理逻辑我们通常会写在 __getitem__ 这个函数内。

数据增强库 Albumentations

在数据集中,加载图像和对应标注(边界框),这里数据增强使用了 Albumentations 这个库,这个库支持 pytorch,适用于目标检测、语义分割等图像增强任务。

import os
import numpy as np
import torch

from PIL import Image,ImageFile
from torch.utils.data import Dataset, DataLoader

这里引入依赖 DatasetDataLoader

import matplotlib.pyplot as plt
from utils import (non_max_suppression as nms, plot_image)

对于图像分类任务,图像(数据)加载比较容易,而对于目标检测的任务,需要对数据进行处理,需要原始数据边界框进行处理为模型可以识别的格式。

COCO128

这里为了演示采用 coco 128 数据集,数据集中有 128 样本,涵盖了 80 类别,主要就是用来检验我们大家网络是否能够正常运行。


import glob
# coco128 数据集

# 选择一个样作为研究对象
sample = 1

# 图像格式
image_formate = "jpg"
# 标注格式
label_formate = "txt"

# 项目目录
project_path = os.getcwd()

# 数据库目录
dataset_path = "data\\coco128\\"

"""
- coco128
-- images
--- train2017
-- labels
--- train2017
"""

# 图像和标注存放位置
img_dir = "images\\train2017"
label_dir = "labels\\train2017"

# 标注类别
category_names = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
        'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow',
        'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
        'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard',
        'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
        'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
        'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone',
        'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
        'hair drier', 'toothbrush'] 

# 图像文件列表
img_paths = glob.glob(os.path.join(project_path,dataset_path,img_dir,"*.jpg"))
# 标注文件列表
label_paths = glob.glob(os.path.join(project_path,dataset_path,label_dir,"*.txt"))
img_path = img_paths[sample]
im = Image.open(img_path)
im

先验边界框的比例

在下面列表(list)中,每一行代表在一个尺度特征图下生成的锚框的尺寸,每个 tuple 表示一个锚框的宽度和高度,这里宽度和高度的值是相对于图像宽度和高度比率。每一组包括 3 不同比例宽高尺寸。例如第一行比例较大比例适合比较粗旷(低分辨率的网格)例如 13×1313 \times 13,这些特征图主要用于检测比较大的目标。

这些锚框的尺寸并不是凭空想象的,而是通过对数据集中所有真实边界框通过聚类算法统计计算得来的。所以这些锚框尺寸是依赖于数据集,不同数据集会有不同比例的锚框尺寸。

anchors = [
    [(0.28,0.22),(0.38,0.48),(0.9,0.78)],
    [(0.07,0.15),(0.15,0.11),(0.14,0.29)],
    [(0.02,0.03),(0.04,0.07),(0.08,0.06)],
]

我们来看 anchors 这个数据结构,这是列表形式,每一行包含 3 个 tuple 格式数据,tuple 数据分别表示锚框的尺寸,这里锚框的宽高尺寸是相对于图像宽度和高度的。第一行锚框的尺寸是对应 13×1313 \times 13 分辨率比较低的特征图,适合预测比较大的目标,而随后依次逐渐减小的锚框尺寸用于预测比较小的目标。

对于不同数据集采用锚框的尺寸是不同的,这个需要通过 kmean 聚类方法来计算锚框的尺寸。聚类过程是采用 1- IoU 作为距离

def iou_width_height(boxes1,boxes2):
    """
    Parameters:
        boxes1(tensor): 第一个边界框的宽度和高度
        boxes2(tensor): 第二个边界框的宽度和高度
    Returns:
        tensor:两个边界框的交并比
    """
    intersection = torch.min(boxes1[...,0],boxes2[...,0]) * torch.min(boxes1[...,1],boxes2[...,1])
    union = (boxes1[...,0]*boxes1[...,1] + boxes2[...,0]*boxes2[...,1] - intersection)
    
    return intersection / union
    
class COCO128Dataset(Dataset):
    def __init__(self, 
                 img_dir,#图像存放目录 
                 label_dir, #标注存放目录
                 anchors, #锚框
                 image_size=416,#图像尺寸
                 S=[13,26,52],#输出图像 3 个尺度
                 C=80,#数据集的类别数
                 transform=None):
        
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.image_size = image_size
        
        self.transform = transform
        self.S = S
        #这里将锚框进行展平为(9,2)也就是对于每个真实边界框都
        self.anchors = torch.tensor(anchors[0] + anchors[1] + anchors[2])
        #锚框的数目
        self.num_anchors = self.anchors.shape[0]
        #特征图输出尺度数(13,26,52)
        self.scale_number = len(S)
        #锚框在每个尺度下锚框数目
        self.num_anchors_per_scale = self.num_anchors // self.scale_number
        #类别数
        self.C = C
        """
        关于正负样本定义,这部分内容对于自己比较含糊,首先正负样本的定义,其次在损失函数中,正负样本如何应用,
        我们现在数据集中标注是一个边界框,所以在数据集中只提供了正例样本,那么应该如何定义负例样本,因为可能
        我们随便在图像框出一框就是负例样本。
        
        对于每一个真实边界框都会分配一个先验边界框,也就是某一个先验边界框与目标 GT 重合最大,为正例样本,
        如果其他先验框也与真实框 IoU 大于指定阈值 0.5 这个先验框既不会当作正例样本,也不会当作负例样本,
        对于没有分配给任何真实 GT 样本,这样样本不会计算类别损失和定位损失。
        """
        self.ignore_iou_thresh = 0.5
        
        self.img_paths = glob.glob(os.path.join(self.img_dir,"*.jpg"))
        self.label_paths = glob.glob(os.path.join(self.label_dir,"*.txt"))
        
    def __len__(self):
        return len(glob.glob(os.path.join(project_path,dataset_path,label_dir,"*.txt")))
    
    def __getitem__(self, index):
        #图像路径和标注文件路径
        image_path = self.img_paths[index]
        label_path = self.label_paths[index]
        #这里使用 roll 就是将 class label 位置调整到末端,[x,y,w,h,class] np.roll 
        #我们来看一看 roll
        bboxes = np.roll(np.loadtxt(fname=label_path,delimiter=" ",ndmin=2),4,axis=1).tolist()
        # 在 getitem 方法中,根据索引来找到图像和图像对应txt文件,然后经过对图像进行变换再输入到
        # 模型,这样可以让模型看到更多样化的图像。这里数据增强采用了 Albumentations 这个库,这个库
        # 要求输入图像和标注数据格式都是 numpy
        image = np.array(Image.open(image_path).convert("RGB"))
        
        #数据增强
        if self.transform:
            augmentations = self.transform(image=image,bboxes=bboxes)
            image = augmentations["image"]
            bboxes = augmentations["bboxes"]
        
        # 这里的 targets 是列表,模型包括多少输出尺度,这里就有多少元素,列表中每个元素对应一个尺度
        # 每一个尺度,也就是每个元素的格式为(3,S,S,6),6 包括目标置信度、位置信息(x,y,w,h) 还有就是
        # 类别信息
        targets = [torch.zeros((self.num_anchors // self.scale_number, S, S, 6)) for S in self.S]
        #print(targets[0].shape)
        #遍历每一个边界框
        for box in bboxes:
            #计算不同尺度上的所有先验框与真实边界框的 IoU 值
            iou_anchors = iou_width_height(torch.tensor(box[2:4]),self.anchors)
            #根据真实边界框和锚框计算 IoU 大小进行排序,按大小排序输出结果为其在原列表该先验框位置
            anchor_indices = iou_anchors.argsort(descending=True,dim=0)
            # 获取真实框的中心点、宽和高以及类别信息
            x,y, width,height, class_label = box
            # 每一个尺度下 anchor
            has_anchor = [False] * 3
            # 从到
            for anchor_idx in anchor_indices:
                #计算出其所在尺度特征图
                scale_idx = anchor_idx // self.num_anchors_per_scale
                #计算在该特征图上第几个锚框
                anchor_on_scale = anchor_idx % self.num_anchors_per_scale
                #获取所在尺度特征图的网格数量
                S = self.S[scale_idx]
                #计算出其所在网格中位置
                i, j  = int(S * y),int(S * x) #网格位置
                #根据边界框所在特征图尺度信息以及所在位置来
                anchor_token = targets[scale_idx][anchor_on_scale,i,j,0]
                # 如果所在位置为 0 并且所在尺度特征图为 False,表示一个边界框在一个尺度特征图的一个网格位置
                # 如果发现该锚点没有被占用,当前边界框在该尺度上还没有任何锚框
                if not anchor_token and not has_anchor[scale_idx]:
                    #更新状态
                    
                    # 就将该锚点分配给真实边界框,并且更新锚点为 1
                    # 表示这个网格存在一个物体,并且更新中心值和宽高
                    targets[scale_idx][anchor_on_scale,i,j,0] =1
                    x_cell, y_cell = S * x - j, S * y -i #在网格心中坐标
                    width_cell, height_cell = (
                        width * S,
                        height * S,
                    )
                    box_coordinates = torch.tensor(
                        [x_cell, y_cell,width_cell,height_cell]
                    )
                    targets[scale_idx][anchor_on_scale, i, j, 1:5] = box_coordinates
                    targets[scale_idx][anchor_on_scale, i, j, 5] = int(class_label)
                    #更新该尺度特征状态,表示在该尺度下已经存在该目标的所对应的锚框
                    has_anchor[scale_idx] = True
                    
                # 如果锚框,该尺度下存在锚框,并且该目标边界框与锚框的 IoU 大于,将值
                elif not anchor_token and iou_anchors[anchor_idx] > self.ignore_iou_thresh:
                    targets[scale_idx][anchor_on_scale,i,j,0] = -1
                    
        return image,tuple(targets)
import albumentations as A
from albumentations.pytorch import ToTensorV2
import cv2
IMAGE_SIZE = 416
scale = 1.1
train_transforms = A.Compose(
    [
        A.LongestMaxSize(max_size=int(IMAGE_SIZE * scale)),
        A.PadIfNeeded(
            min_height=int(IMAGE_SIZE * scale),
            min_width=int(IMAGE_SIZE * scale),
            border_mode=cv2.BORDER_CONSTANT,
        ),
        A.RandomCrop(width=IMAGE_SIZE,height=IMAGE_SIZE),
        A.ColorJitter(brightness=0.6,contrast=0.6,saturation=0.6,hue=0.6,p=0.4),
        A.OneOf(
            [
                A.ShiftScaleRotate(
                    rotate_limit=10,p=0.4,border_mode=cv2.BORDER_CONSTANT
                ),
#                 A.IAAAffine(shear=10,p=0.4,model="constant"),
            ],
            p=1.0
        ),
        A.HorizontalFlip(p=0.5),
        A.Blur(p=0.1),
        A.CLAHE(p=0.1),
        A.Posterize(p=0.1),
        A.ToGray(p=0.1),
        A.ChannelShuffle(p=0.05),
        A.Normalize(mean=[0,0,0],std=[1,1,1],max_pixel_value=255,),
        ToTensorV2(),
    ],
    bbox_params=A.BboxParams(format="yolo",min_visibility=0.4,label_fields=[],),
)
def test():
    # 定义数据集
    dataset = COCO128Dataset(
        img_dir = os.path.join(project_path,dataset_path,img_dir),
        label_dir = os.path.join(project_path,dataset_path,label_dir),
        anchors = anchors,
        transform = train_transforms)
    #预测输出 3 种不同类型尺度
    S = [13,26,52]
    
    scaled_anchors = torch.tensor(anchors)/ (
        1 / torch.tensor(S).unsqueeze(1).unsqueeze(1).repeat(1,3,2)
    )
    loader = DataLoader(dataset=dataset,batch_size=1,shuffle=True)
    # x 表示图像,y 表示标注
    for x, y in loader:
        boxes = []
        
        for i in range(y[0].shape[1]):
            anchor = scaled_anchors[i]
            boxes += cells_to_bboxes(
                y[i],is_preds=False,S=y[i].shape[2],anchors=anchors
            )[0]
        boxes = nms(boxes,iou_threshold=1,prob_threshold=0.7,box_format="midpoint")
        plot_image(x[0].permute(1,2,0),boxes)

test()

当我们加载该图像的标注时,标注是一个边界框的数组。随后在训练过程通过损失函数来计算真实边界框和预测边界框之间差值,通过缩小这个误差来更新模型的参数。这样一来就需要将模型预测的边界框和标注的真实边界框在一个尺度下(也就是可对比),模型会从 3 种不同尺度上输出预测边界框,也就是在不同尺度的特征图分别创建一个真实边界框(或者随后也叫做目标边界框,都是一个意思)

其中 6 表示目标置信度得分、边界框的 x,y,w,h 坐标以及类别数,这里做了前提假设,也就是每个边界框只有一个标签,而且在每一个尺度上边界框的数量是相等,targets = [torch.zeros((self.num_anchors // self.scale_number, S, S, 6)) for S in self.S] 初始化了 3 不同目标 tensors ,初始值都是为 0 S 不同网格大小,从小的大依次为 13、26 和 52

def cells_to_bboxes(preds, anchors, S, is_preds=True):
    """
    
    """
    #获取 batch_size
    BATCH_SIZE = preds.shape[0]
    # 锚框数量
    num_anchors = len(anchors)
    #获取预测边界框的 x,y w,h 坐标
    box_preds = preds[...,1:5]
    # 是否为预测
    if is_preds:
        # reshape
        anchors = anchors.reshape(1,len(anchors),1,1,2)
        box_preds[...,0:2] = torch.sigmoid(box_preds[...,0:2])
        box_preds[...,2:] = torch.exp(box_preds[...,2:]) * anchors
        scores = torch.sigmoid(preds[...,0:1])
        best_class = torch.argmax(preds[...,5:],dim=-1).unsqueeze(-1)
    else:
        #获取置信度
        scores = preds[...,0:1]
        best_class = preds[...,5:6]
        
    cell_indices = (
        torch.arange(S)
        # batch_size, scale_number grid size 
        .repeat(preds.shape[0],3,S,1)
        .unsqueeze(-1)
    )
    
    x = 1 / S * (box_preds[...,0:1] + cell_indices)
    y = 1 / S * (box_preds[...,1:2] + cell_indices.permute(0,1,3,2,4))
    w_h = 1 / S * box_preds[...,2:4]
    
    converted_bboxes = torch.cat((best_class,scores,x,y,w_h),dim=-1).reshape(BATCH_SIZE,num_anchors*S*S,6)
    return converted_bboxes.tolist()