yolov5-5.0版本代码详解----datasets.py的LoadImagesAndLabels函数（2）

yolov5-5.0版本代码详解----datasets.py的LoadImagesAndLabels函数（2）

1、__len__函数

类的特殊属性，通过重写他，能让内置函数len()的参数是自定义类型

    def __len__(self):
        return len(self.img_files)  # 图片的全路径列表

2、__getitem__函数

    # 如果在类中定义了__getitem__()方法，
    # 那么他的实例对象（假设为P）就可以这样P[key]取值。
    # 当实例对象做P[key]运算时，就会调用类中的__getitem__()方法
    # 返回__getitem__的return
    # 例如在__getiten__ return 123
    # ni = 类（）
    # print(ni)
    # 就都会 输出：123

2.1 获取要进行数据增强的图片index、hyp等

        index = self.indices[index]  # self.indices = range(n)  # 所有图片的index
        hyp = self.hyp  # 超参 包含众多数据增强超参
        mosaic = self.mosaic and random.random() < hyp['mosaic']  # 是否进行马赛克数据增强

2.2 mosaic + MixUp

      if mosaic:
            img, labels = load_mosaic(self, index)
            shapes = None

            # 查看马赛克结果图
            # cv2.imshow("masaic", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # mixup数据增强
            if random.random() < hyp['mixup']:  # 0则关闭 1则100%打开
                # *load_mosaic(self, random.randint(0, self.n - 1)) 随机从数据集中任选一张图片和本张图片进行mixup数据增强
                # img:   两张图片融合之后的图片--numpy (640, 640, 3)
                # labels: 两张图片融合之后的标签--label [M+N, cls+x1y1x2y2]
                img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1)))

                # 测试代码 测试MixUp效果
                # cv2.imshow("MixUp", img)
                # cv2.waitKey(0)
                # cv2.destroyAllWindows()

2.3 不进行马赛克、进行load_image与Letterbox

        else:
            # 3、载入图片
            # 载入图片后还会进行一次resize  将当前图片的最长边缩放到指定的大小, 较小边同比例缩放
            # return：
            # 1、img: resize后的图片；   2、(h0, w0): 原始图片的hw； 3、(h, w): resize后的图片的hw
            img, (h0, w0), (h, w) = load_image(self, index)

            # 测试load_image效果
            # cv2.imshow("load_image", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # 4、Letterbox 主要运用于val或者detect
            # 4.1 确定这张当前图片letterbox之后的shape
            # 如果不用self.rect矩形训练shape就是self.img_size
            shape = self.batch_shapes[self.batch[index]] if self.rect else self.img_size

            # 4.2 letterbox 这一步将第一步缩放得到的图片再缩放到当前batch所需要的尺度
            # 矩形推理需要一个batch的所有图片的shape必须相同
            # 这里没有缩放操作，所以这里的ratio永远都是(1.0, 1.0)  pad=(0.0, 20.5)
            img, ratio, pad = letterbox(img, shape, auto=False, scaleup=self.augment)
            shapes = (h0, w0), ((h / h0, w / w0), pad)

            # 图片letterbox之后label的坐标也要相应变化  根据pad调整label坐标 并将归一化的xywh -> 未归一化的xyxy
            labels = self.labels[index].copy()
            if labels.size:
                labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1])

            # 测试代码 测试letterbox效果
            # cv2.imshow("letterbox", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

2.5 random_perspective增强

            if self.augment:
                # 5、random_perspective增强: 随机对图片进行旋转，平移，缩放，裁剪，透视变换
                img, labels = random_perspective(img, labels,
                                                 degrees=hyp['degrees'],
                                                 translate=hyp['translate'],
                                                 scale=hyp['scale'],
                                                 shear=hyp['shear'],
                                                 perspective=hyp['perspective'])

        nl = len(labels)  # labels的数量
        if nl:  # 若不为0
            # 转换label尺寸格式 将未归一化的xyxy -> 归一化的xywh
            labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3)

2.6 平移增强 随机左右翻转 + 随机上下翻转

        if self.augment:
            img, labels = self.albumentations(img, labels)

            # 色域空间增强
            augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v'])
            # 测试代码 测试augment_hsv效果
            # cv2.imshow("augment_hsv", img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

            # 6.1 随机上下翻转 flip up-down
            if random.random() < hyp['flipud']:
                img = np.flipud(img)  # np.flipud 将数组在上下方向翻转。
                if nl:
                    labels[:, 2] = 1 - labels[:, 2]  # 1 - y_center  label也要映射

            # 6.2 随机左右翻转 flip left-right
            if random.random() < hyp['fliplr']:
                img = np.fliplr(img)  # np.fliplr 将数组在左右方向翻转
                if nl:
                    labels[:, 1] = 1 - labels[:, 1]  # 1 - x_center  label也要映射

2.7 初始化标签框对应的图片序号, 配合下面的collate_fn使用

        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
        img = np.ascontiguousarray(img)  # 看不懂 img变成内存连续的数据  加快运算

2.8 返回值

return torch.from_numpy(img), labels_out, self.img_files[index], shapes
"""   
返回值：
1、torch.from_numpy(img): 这个index的图片数据(增强后) [3, 640, 640]
2、labels_out: 这个index图片的gt label [6, 6] = [gt_num, 0+class+xywh(normalized)]
3、self.img_files[index]: 这个index图片的路径地址
4、shapes: 这个batch的图片的shapes 测试时(矩形训练)才有  验证时为None
"""

3、collate_fn函数

    # 在create_dataloader中生成dataloader时调用
    def collate_fn(batch):
        """
        整理函数  将image和label整合到一起
        pytorch的DataLoader打包一个batch的数据集时要经过此函数进行打包 通过重写此函数实现标签与图片对应的划分，
        一个batch中哪些标签属于哪一张图片,形如
            [[0, 6, 0.5, 0.5, 0.26, 0.35],
             [0, 6, 0.5, 0.5, 0.26, 0.35],
             [1, 6, 0.5, 0.5, 0.26, 0.35],
             [2, 6, 0.5, 0.5, 0.26, 0.35],]
             图片编号、labels类别编号 xywh
        """
        # img: 一个tuple 由batch_size个tensor组成 整个batch中每个tensor表示一张图片
        # label: 一个tuple 由batch_size个tensor组成 每个tensor存放一张图片的所有的target信息
        # label[6, object_num] 6中的第一个数代表一个batch中的第几张图
        # path: 一个tuple 由4个str组成, 每个str对应一张图片的地址信息
        img, label, path, shapes = zip(*batch)  # transposed
        for i, l in enumerate(label):
            l[:, 0] = i  
        
        # torch.stack(img, 0): 整个batch的图片--numpy
        # torch.cat(label, 0): [num_target, img_index+class_index+xywh(normalized)] 整个batch的label
        # path: 整个batch所有图片的路径
        # shapes: (h0, w0), ((h / h0, w / w0), pad)
        return torch.stack(img, 0), torch.cat(label, 0), path, shapes

4、collate_fn4函数

    def collate_fn4(batch):
        """同样在create_dataloader中生成dataloader时调用：
        这里是yolo-v5作者实验性的一个代码 quad-collate function 当train.py的opt参数quad=True 则调用collate_fn4代替collate_fn
        作用:  如之前用collate_fn可以返回图片[16, 3, 640, 640] 经过collate_fn4则返回图片[4, 3, 1280, 1280]
              将4张mosaic图片[1, 3, 640, 640]合成一张大的mosaic图片[1, 3, 1280, 1280]
              将一个batch的图片每四张处理, 0.5的概率将四张图片拼接到一张大图上训练, 0.5概率直接将某张图片上采样两倍训练
        """
        img, label, path, shapes = zip(*batch)  # transposed
        n = len(shapes) // 4
        img4, label4, path4, shapes4 = [], [], path[:n], shapes[:n]

        ho = torch.tensor([[0., 0, 0, 1, 0, 0]])
        wo = torch.tensor([[0., 0, 1, 0, 0, 0]])
        s = torch.tensor([[1, 1, .5, .5, .5, .5]])  # scale
        for i in range(n):  # zidane torch.zeros(16,3,720,1280)  # BCHW
            i *= 4
            if random.random() < 0.5:
                im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2., mode='bilinear', align_corners=False)[
                    0].type(img[i].type())
                l = label[i]
            else:
                im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2)
                l = torch.cat((label[i], label[i + 1] + ho, label[i + 2] + wo, label[i + 3] + ho + wo), 0) * s
            img4.append(im)
            label4.append(l)

        for i, l in enumerate(label4):
            l[:, 0] = i  # add target image index for build_targets()

        return torch.stack(img4, 0), torch.cat(label4, 0), path4, shapes4