yolov5代码阅读笔记2

yolov5的dataset和dataloader部分
看这一部分之前，可以先看看之前的文章从零开始dataset

在train.py文件的300行左右可以找到下面这段代码

    # Trainloader
    train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls,
                                              hyp=hyp, augment=True, cache=None if opt.cache == 'val' else opt.cache,
                                              rect=opt.rect, rank=LOCAL_RANK, workers=workers,
                                              image_weights=opt.image_weights, quad=opt.quad,
                                              prefix=colorstr('train: '), shuffle=True)

进入到create_dataloader我们可以看到下面的代码，即创建dataset

        dataset = LoadImagesAndLabels(path, imgsz, batch_size,
                                      augment=augment,  # augmentation
                                      hyp=hyp,  # hyperparameters
                                      rect=rect,  # rectangular batches
                                      cache_images=cache,
                                      single_cls=single_cls,
                                      stride=int(stride),
                                      pad=pad,
                                      image_weights=image_weights,
                                      prefix=prefix)

同样为了阅读代码我们使用下面的代码:

from utils.datasets import LoadImagesAndLabels
import numpy as np
path='/datasets/coco128/images/train2017'
dataset = LoadImagesAndLabels(path)

dataset

我们知道dataset中__init__初始化,__len__数据长度,getitem__得到某个数据
先看__init，init函数就是将图片路径和label路径读取了出来并检查有没有错误，并且又从cache中读取了一次，检查有无错误，然后判断有无缓存，因为有缓存读缓存更快，很多个人独代码可以略过，只需要把图片和label读入就可以了。

    def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False,
                 cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''):
        # 图片输入大小
        self.img_size = img_size
        self.augment = augment#是否开始数据增强
        self.hyp = hyp#超参数
        self.image_weights = image_weights
        self.rect = False if image_weights else rect
        self.mosaic = self.augment and not self.rect  # load 4 images at a time into a mosaic (only during training)
        self.mosaic_border = [-img_size // 2, -img_size // 2]
        self.stride = stride#布长
        self.path = path#训练集路径
        self.albumentations = Albumentations() if augment else None

        try:
            f = []  # image files
            for p in path if isinstance(path, list) else [path]:
                p = Path(p)  # os-agnostic
                # path->p /Users/lrg/Downloads/datasets/coco128/images/train2017
                print(f'path->p:{p}')
                if p.is_dir():  # dir
                    # 添加数据集的所有图片路径
                    # f就是所有图片的路径集合
                    f += glob.glob(str(p / '**' / '*.*'), recursive=True)
                    # f = list(p.rglob('*.*'))  # pathlib
                elif p.is_file():  # file
                    with open(p) as t:
                        t = t.read().strip().splitlines()
                        parent = str(p.parent) + os.sep
                        f += [x.replace('./', parent) if x.startswith('./') else x for x in t]  # local to global path
                        # f += [p.parent / x.lstrip(os.sep) for x in t]  # local to global path (pathlib)
                else:
                    raise Exception(f'{prefix}{p} does not exist')

            '''
                图片的各种格式
                IMG_FORMATS:('bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp')
                self.im_files:数据集图片路径集合
            '''
            self.im_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS)
            # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS])  # pathlib
            assert self.im_files, f'{prefix}No images found'
        except Exception as e:
            raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}')

        # Check cache
        # 得到label的路径地址
        self.label_files = img2label_paths(self.im_files)  # labels
        # 得到cache缓存的地址,关于.cache，可以看1.1
        cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache')
        try:
            # np.load加载.cache数据
            cache, exists = np.load(cache_path, allow_pickle=True).item(), True  # load dict
            assert cache['version'] == self.cache_version  # same version
            assert cache['hash'] == get_hash(self.label_files + self.im_files)  # same hash
        except Exception:
            cache, exists = self.cache_labels(cache_path, prefix), False  # cache
        #想看cache里面内容，可以看1.1
        # Display cache
        nf, nm, ne, nc, n = cache.pop('results')  # found, missing, empty, corrupt, total
        # found:128, missing:0, empty:2, corrupt:0, total:128
        print(f'found:{nf}, missing:{nm}, empty:{ne}, corrupt:{nc}, total:{n}')
        if exists:
            d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupt"
            tqdm(None, desc=prefix + d, total=n, initial=n)  # display cache results
            if cache['msgs']:
                LOGGER.info('\n'.join(cache['msgs']))  # display warnings
        assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}'

        # Read cache
        [cache.pop(k) for k in ('hash', 'version', 'msgs')]  # remove items
        # *cache:取出所有的key
        # *cache.values:取出所有的values
        # labels:128   shapes:128张图片的长宽,是否增强
        labels, shapes, self.segments = zip(*cache.values())
        self.labels = list(labels)
        self.shapes = np.array(shapes, dtype=np.float64)
        # cache中也有图片和label的路径，但是之前代码也读取了相应的路径，所以这儿做了个更新
        self.im_files = list(cache.keys())  # update
        self.label_files = img2label_paths(cache.keys())  # update
        n = len(shapes)  # number of images
        # np.floor()返回不大于输入参数的最大整数
        # bi用来确定某张图片属于第几个batch_size
        bi = np.floor(np.arange(n) / batch_size).astype(np.int)  # batch index
        nb = bi[-1] + 1  # number of batches
        self.batch = bi  # batch index of image
        self.n = n
        self.indices = range(n)

        # Update labels
        # 查看label是否错误
        include_class = []  # filter labels to include only these classes (optional)
        include_class_array = np.array(include_class).reshape(1, -1)
        for i, (label, segment) in enumerate(zip(self.labels, self.segments)):
            # print(f'include_class:{i,include_class}:single_cls:{single_cls}')
            # include_class=[],single_cls=False
            if include_class:
                j = (label[:, 0:1] == include_class_array).any(1)
                self.labels[i] = label[j]
                if segment:
                    self.segments[i] = segment[j]
            if single_cls:  # single-class training, merge all classes into 0
                self.labels[i][:, 0] = 0
                if segment:
                    self.segments[i][:, 0] = 0

        # False,可以不用管，用不到
        # Rectangular Training
        if self.rect:
            # Sort by aspect ratio
            s = self.shapes  # wh
            ar = s[:, 1] / s[:, 0]  # aspect ratio
            irect = ar.argsort()
            self.im_files = [self.im_files[i] for i in irect]
            self.label_files = [self.label_files[i] for i in irect]
            self.labels = [self.labels[i] for i in irect]
            self.shapes = s[irect]  # wh
            ar = ar[irect]

            # Set training image shapes
            shapes = [[1, 1]] * nb
            for i in range(nb):
                ari = ar[bi == i]
                mini, maxi = ari.min(), ari.max()
                if maxi < 1:
                    shapes[i] = [maxi, 1]
                elif mini > 1:
                    shapes[i] = [1, 1 / mini]

            self.batch_shapes = np.ceil(np.array(shapes) * img_size / stride + pad).astype(np.int) * stride

        # Cache images into RAM/disk for faster training (WARNING: large datasets may exceed system resources)
        self.ims = [None] * n
        self.npy_files = [Path(f).with_suffix('.npy') for f in self.im_files]
        print(f'self.npy:{self.npy_files}')
        print('---------------------')
        print(f'what is cache_images{cache_images}')
        # 有缓存图片就直接从disk读取
        if cache_images:
            gb = 0  # Gigabytes of cached images
            self.im_hw0, self.im_hw = [None] * n, [None] * n
            fcn = self.cache_images_to_disk if cache_images == 'disk' else self.load_image
            results = ThreadPool(NUM_THREADS).imap(fcn, range(n))
            pbar = tqdm(enumerate(results), total=n)
            for i, x in pbar:
                if cache_images == 'disk':
                    gb += self.npy_files[i].stat().st_size
                else:  # 'ram'
                    self.ims[i], self.im_hw0[i], self.im_hw[i] = x  # im, hw_orig, hw_resized = load_image(self, i)
                    gb += self.ims[i].nbytes
                pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})'
            pbar.close()

len

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

getitem

    def __getitem__(self, index):
        index = self.indices[index]  # linear, shuffled, or image_weights

        hyp = self.hyp
        # 马赛克增强，
        mosaic = self.mosaic and random.random() < hyp['mosaic']
        # mosaic增强
        # 马赛克数据增强：
        # 可以增加数据的复杂度，同时可以增加一张图片的标注数量
        if mosaic:
            # Load mosaic 见1.2
            img, labels = self.load_mosaic(index)
            shapes = None

            # MixUp augmentation
            # mixup是一种data augmentation方法，可以用来提升模型的泛化能力和对对抗样本(adversarial examples，指的是训练样本分布外的样本)的鲁棒性。
            if random.random() < hyp['mixup']:
                img, labels = mixup(img, labels, *self.load_mosaic(random.randint(0, self.n - 1)))

        else:
            # Load image
            # im, hw_original, hw_resized
            img, (h0, w0), (h, w) = self.load_image(index)

            # Letterbox
            shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size  # final letterboxed shape
            # index: 1, self.img_size=shape: 640
            # letterbox：加灰条  见1.3
            img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
            shapes = (h0, w0), ((h / h0, w / w0), pad)  # for COCO mAP rescaling

            labels = self.labels[index].copy()
            print(f'labels.size:{labels.size}')
            if labels.size:  # normalized xywh to pixel xyxy format
                labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])

            if self.augment:
                '''
                    透视变换
                    yolov5的数据增强中，透视、仿射变换统一使用了random_perspective一个函数进行处理，
                    包含了旋转、缩放、平移、剪切变换
                    对于一张图片，可以使用矩阵对其进行旋转，平移，等操作
                    见1.4
                '''
                img, labels = random_perspective(img, labels,
                                                 degrees=hyp['degrees'],
                                                 translate=hyp['translate'],
                                                 scale=hyp['scale'],
                                                 shear=hyp['shear'],
                                                 perspective=hyp['perspective'])
        '''
            mosaic的数据是xyxy
            img类型是xywh
        '''
        nl = len(labels)  # number of labels
        if nl:
            labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)

        # 下面就是HSV增强，上下旋转等，之前都介绍过
        if self.augment:
            # Albumentations
            # albumentations 是一个给予 OpenCV的快速训练数据增强库，拥有非常简单且强大的可以用于多种任务（分割、检测）的接口，易于定制且添加其他框架非常方便。
            # 它可以对数据集进行逐像素的转换，如模糊、下采样、高斯造点、高斯模糊、动态模糊、RGB转换、随机雾化等；也可以进行空间转换（同时也会对目标进行转换），如裁剪、翻转、随机裁剪等。
            img, labels = self.albumentations(img, labels)
            nl = len(labels)  # update after albumentations

            # HSV color-space
            augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])

            # Flip up-down
            if random.random() < hyp['flipud']:
                img = np.flipud(img)
                if nl:
                    labels[:, 2] = 1 - labels[:, 2]

            # Flip left-right
            if random.random() < hyp['fliplr']:
                img = np.fliplr(img)
                if nl:
                    labels[:, 1] = 1 - labels[:, 1]

            # Cutouts
            # labels = cutout(img, labels, p=0.5)
            # nl = len(labels)  # update after cutout

        labels_out = torch.zeros((nl, 6))
        if nl:
            labels_out[:, 1:] = torch.from_numpy(labels)

        # Convert
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        # ascontiguousarray函数将一个内存不连续存储的数组转换为内存连续存储的数组，使得运行速度更快
        img = np.ascontiguousarray(img)

        return torch.from_numpy(img), labels_out, self.im_files[index], shapes

1.1 cache

cache_path='/Users/lrg/Downloads/datasets/coco128/labels/train2017.cache'
cache=np.load(cache_path,allow_pickle=True).item()
print(type(cache))
for k in cache.keys():
    try:
    	#输出图片信息
        if k.find('.jpg')!=-1:
            print(f'key:{k},,value:{cache[k][0].shape},size:{cache[k][1]},[]:{cache[k][2]}')
        else:
        	#输出额外信息
            print(f'key:{k},values:{cache[k]}')
    except:
        print(k.find('.jpg'),k,cache[k][0])

打印出来的信息

1.2 Load mosaic

mosaic增强马赛克数据增强：可以增加数据的复杂度，同时可以增加一张图片的标注数量,就是把四张图片拼接在一起合成一张图片

import cv2
import numpy as np
import random
from PIL import Image
cache_path='/Users/lrg/Downloads/datasets/coco128/labels/train2017.cache'
cache=np.load(cache_path,allow_pickle=True).item()

cache.pop('results')
[cache.pop(k) for k in ('hash', 'version', 'msgs')]  # remove items
# segments=[]
labels, shapes, segmentss = zip(*cache.values())
labelss = list(labels)
def load_mosaic():
    # YOLOv5 4-mosaic loader. Loads 1 image + 3 random images into a 4-image mosaic
    labels4, segments4 = [], []
    s = 640
    mosaic_border=[-320,-320]
    # random.uniform(x, y)方法将随机生成一个实数，它在 [x,y] 范围内。
    # [-x,2*s+x]-->[320,960]
    # 随机生成一个点作为四张图片的拼接的中心点xc,yc
    x=-320
    index=1
    print(-x,2*s+x)
    indices=range(128)
    yc, xc = (int(random.uniform(-x, 2*s+x)) for x in mosaic_border)  # mosaic center x, y
    print(xc,yc)
    # 从非空序列中随机选取一个数据并返回
    # 随机选择3个数字和index组成4个   ，合成一幅图片
    indices = [index] + random.choices(indices, k=3)  # 3 additional image indices
    print(indices)
    # 打乱顺序
    random.shuffle(indices)
    print(indices)
    # 自己测试，使用了同一张图片
    namelist=['bus.jpg','bus2.jpg','bus.jpg','bus2.jpg']
    for i, index in enumerate(namelist):
        # Load image
        print(i,index)
        img= cv2.imread(index)
        h,w=img.shape[:2]
        '''
            np.full 构造一个数组，用指定值填充其元素
            full(shape, fill_value, dtype=None, order='C')
            shape：int 或者 int元组
            fill_value：填充到数组中的值
            np.full：制作一幅能容纳4张图的图片img4，尺寸为(2*s.2*s,3),先用114(灰色)填充
            然后把4幅图复制到img4上
            
        '''
        # place img in img4
        # 以xc,yc为中心点，将四张图片拼接
        if i == 0:  # top left看下图
            img4 = np.full((s * 2, s * 2, img.shape[2]), 114, dtype=np.uint8)  # base image with 4 tiles
            x1a, y1a, x2a, y2a = max(xc - w, 0), max(yc - h, 0), xc, yc  # xmin, ymin, xmax, ymax (large image)
            x1b, y1b, x2b, y2b = w - (x2a - x1a), h - (y2a - y1a), w, h  # xmin, ymin, xmax, ymax (small image)
        elif i == 1:  # top right
            x1a, y1a, x2a, y2a = xc, max(yc - h, 0), min(xc + w, s * 2), yc
            x1b, y1b, x2b, y2b = 0, h - (y2a - y1a), min(w, x2a - x1a), h
        elif i == 2:  # bottom left
            x1a, y1a, x2a, y2a = max(xc - w, 0), yc, xc, min(s * 2, yc + h)
            x1b, y1b, x2b, y2b = w - (x2a - x1a), 0, w, min(y2a - y1a, h)
        elif i == 3:  # bottom right
            x1a, y1a, x2a, y2a = xc, yc, min(xc + w, s * 2), min(s * 2, yc + h)
            x1b, y1b, x2b, y2b = 0, 0, min(w, x2a - x1a), min(y2a - y1a, h)
    #
        img4[y1a:y2a, x1a:x2a] = img[y1b:y2b, x1b:x2b]  # img4[ymin:ymax, xmin:xmax]
        padw = x1a - x1b
        padh = y1a - y1b

    #
        global labelss
        global segmentss
    #     # Labels
    #     因为改动了循环代码，所以index不再指indices的循环内容，所以自己手动改回来
        labels, segments = labelss[indices[i]].copy(), segmentss[indices[i]].copy()
        # if labels.size:
        #     labels[:, 1:] = xywhn2xyxy(labels[:, 1:], w, h, padw, padh)  # normalized xywh to pixel xyxy format
        #     segments = [xyn2xy(x, w, h, padw, padh) for x in segments]
        labels4.append(labels)
        segments4.extend(segments)
    im = Image.fromarray(img4)
    # im.show()
    # Concat/clip labels
    # 可以自己打印一下看shape变化 --变为了(n,5)
    print(len(labels4))
    labels4 = np.concatenate(labels4, 0)
    print(labels4.shape)
    # np.clip是一个截取函数，用于截取数组中小于或者大于某值的部分，并使得被截取部分等于固定值
    for x in (labels4[:, 1:], *segments4):
        np.clip(x, 0, 2 * s, out=x)  # clip when using random_perspective()

    # Augment
    # img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste'])
    # 透视变化见1.3
    # img4, labels4 = random_perspective(img4, labels4, segments4,
    #                                    degrees=self.hyp['degrees'],
    #                                    translate=self.hyp['translate'],
    #                                    scale=self.hyp['scale'],
    #                                    shear=self.hyp['shear'],
    #                                    perspective=self.hyp['perspective'],
    #                                    border=self.mosaic_border)  # border to remove
    #
    # return img4, labels4
load_mosaic()

假装此处有图

1.3 letterbox

基本原理与之前的缩放没什么区别

import cv2
import numpy as np
path='/Users/lrg/Downloads/datasets/coco128/images/train2017/000000000009.jpg'
im = cv2.imread(path)  # BGR
shape = im.shape[:2]  # current shape [height, width]
print(f'current:{shape}')
new_shape=(320, 320)
color=(114, 114, 114)#灰条

# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
scaleup=True
if not scaleup:  # only scale down, do not scale up (for better val mAP)
    r = min(r, 1.0)
ratio=r,r
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
# (640, 480)
print(new_unpad)
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
print(f'之前的dw:{dw},dh:{dh}')
auto=True
scaleFill=False
stride=32
# 此处取模是为了刚好缩放后的图片尺寸能是stride的倍数
if auto:  # minimum rectangle
    # 相当于Python模运算符``x1％x2``，并且与除数x2具有相同的符号
    dw, dh = np.mod(dw, stride), np.mod(dh, stride)  # wh padding
    print(f'此时的dw:{dw},dh:{dh}')
elif scaleFill:  # stretch
    # 铺满
    dw, dh = 0.0, 0.0
    new_unpad = (new_shape[1], new_shape[0])
    ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios
'''
    为什么要取模
'''
dw /= 2  # divide padding into 2 sides
dh /= 2
#
if shape[::-1] != new_unpad:  # resize
    im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
print(top,bottom,left,right)
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
import matplotlib.pyplot as plt
plt.imshow(im)
plt.show()

1.4 random_perspective

只是看了一下旋转的操作,具体可看这篇文章random_perspective

import numpy as np
import cv2
import math
import random
perspective=0.0
degrees=10
translate=.1
scale=.1
shear=10
border=(0, 0)
path='../ATest/bus.jpg'
im=cv2.imread(path)
height = im.shape[0] + border[0] * 2  # shape(h,w,c)
width = im.shape[1] + border[1] * 2

# Center   3x3单位矩阵
C = np.eye(3)
C[0, 2] = -im.shape[1] / 2  # x translation (pixels)
C[1, 2] = -im.shape[0] / 2  # y translation (pixels)

# Perspective 透视变换
P = np.eye(3)
P[2, 0] = random.uniform(-perspective, perspective)  # x perspective (about y)
P[2, 1] = random.uniform(-perspective, perspective)  # y perspective (about x)

# Rotation and Scale # 设置旋转和缩放的仿射矩阵
R = np.eye(3)
a = random.uniform(-degrees, degrees)
# a += random.choice([-180, -90, 0, 90])  # add 90deg rotations to small rotations
s = random.uniform(1 - scale, 1 + scale)
# s = 2 ** random.uniform(-scale, scale)
'''
M=cv2.getRotationMatrix2D(center, angle, scale)
center表示中间点的位置，angle表示旋转5度，scale表示进行等比列的缩放
'''
R[:2] = cv2.getRotationMatrix2D(angle=a, center=(0, 0), scale=s)

# Shear；设置裁剪的仿射矩阵系数
S = np.eye(3)
S[0, 1] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # x shear (deg)
S[1, 0] = math.tan(random.uniform(-shear, shear) * math.pi / 180)  # y shear (deg)

# Translation；设置平移的仿射矩阵系数
T = np.eye(3)
T[0, 2] = random.uniform(0.5 - translate, 0.5 + translate) * width  # x translation (pixels)
T[1, 2] = random.uniform(0.5 - translate, 0.5 + translate) * height  # y translation (pixels)
# 融合仿射矩阵并作用在图片上； @表示矩阵乘法运算
# Combined rotation matrix
# M = T @ S @ R @ P @ C  # order of operations (right to left) is IMPORTANT
M=R#只是旋转图片矩阵
print(M)
# print(perspective)
# print(border)
if (border[0] != 0) or (border[1] != 0) or (M != np.eye(3)).any():  # image changed
    '''
        cv2.warpAffine()放射变换函数，可实现旋转，平移，缩放；变换后的平行线依旧平行
        warpAffine ：意思是仿射变化
    '''
    if perspective:
        im2 = cv2.warpPerspective(im, M, dsize=(width, height), borderValue=(114, 114, 114))
    else:  # affine
        im2 = cv2.warpAffine(im, M[:2], dsize=(width, height), borderValue=(114, 114, 114))

# Visualize
import matplotlib.pyplot as plt
ax = plt.subplots(1, 2, figsize=(12, 6))[1].ravel()
ax[0].imshow(im[:, :, ::-1])  # base
ax[1].imshow(im2[:, :, ::-1])  # warped
plt.show()

可以得到输出的旋转矩阵
[[ 1.00634312 -0.14747734 0. ]
[ 0.14747734 1.00634312 -0. ]
[ 0. 0. 1. ]]
旋转后的图片

dataloader

最后返回dataloader

    return loader(dataset,
                  batch_size=batch_size,
                  shuffle=shuffle and sampler is None,
                  num_workers=nw,
                  sampler=sampler,
                  pin_memory=True,
                  collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn), dataset