YOLO系列----（一）最基础也是最重要的训练数据读取

想要训练一个YOLO的网络，数据集的准备是必不可缺的。初学深度学习的时候，使用的都是官方的数据集。但实际场景中会根据不同的需要使用不同的标注软件来进行标注，所以就需要自己来制作数据集。

想要自己制作数据集，那就要明白数据在训练时会用到哪些信息？

在分类的时候只需要两个信息：图片（image）和类别（label）

而在目标检测中，我们除了要知道图片的类别以外，还需要知道物体的位置。也就是需要三个信息：图片（image）、标注框的坐标（boxes）、框对应的类别（label）

1. 在正式开始训练前会生成一个txt文件，文件里面是存储用来训练的图片以及每张图片中物体对应的框的信息，一行代表一张图片的信息。这里不同的信息之间是用空格来分开的，以便于后面处理。【第一个表示图片的路径，后面的就是每个标注框box的坐标信息(左、上、右、下），box这个坐标是相对于原图的】。train.txt的格式是这样子的：

2.有了上面的这些信息，需要将这些信息转换成计算机可以识别的东西。原始时候是一个txt文件，里面可能有很多行，所以需要对每一行做一次数据处理操作，这样就得到了每一张图片的图片（image）、标注框的坐标（boxes）、框对应的类别（label）信息。

3.训练的时候很少会有人直接用原图来训练，而且在在YOLO中训练时图片的输入大小必须是32的倍数，所以在制作自己的数据集时肯定会对图片做一些前处理操作。在这里就是直接对图片进行resize（改变图像大小）操作，相应的也要对box对进行调整，使得box是相对于resize之后的图片

类似的，在做数据增强的时候如果图像（大小、方向）发生了改动，box也要做相对应的改动

Last:代码实现部分

代码实现获取图片（image）、标注框的坐标（boxes）、框对应的类别（label）是在_getitem_这个函数里面实现的，但是你会发现返回的时候还返回了三个结果image, boxes,y_true

image就是读取到的图片信息，boxes是包含每个框以及框的类别信息，y_true是在训练过程中求loss的时候会用到的一个参数，这里可以先不考虑这个参数（y_true也是较难理解的一部分）。

在定义GetDatasets这个类的时候，annotation_line这个变量就是读取txt得到的列表。

import cv2
import random
import numpy as np
from torch.utils.data.dataset import Dataset



class GetDatasets(Dataset):
    # 1.首先要了解搭建自己的数据集需要什么？
    # 答：(1)数据的路径   （2）resize图片的大小   (3）种类的个数   (4)anchors
    def __init__(self, annotation_line, resize_hw, numm_classes, anchors, is_augmen):
        super(GetDatasets, self).__init__()
        self.anchors = anchors
        self.resize_hw = resize_hw
        self.num_cls = numm_classes
        self.data_augmen = is_augmen
        self.anno_line = annotation_line
        self.bbox_attrs = 5 + numm_classes
        self.anchors_mask = [[6, 7, 8], [3, 4, 5], [0, 1, 2]]

    def __getitem__(self, index):
        # 2.了解自己读取数据最终是获取到什么信息
        # 答: (1)图片image   （2）坐标和种类信息：[x1,y1,x2,y2,cls]===>左上角:(x1,y1) |  右下角:(x2,y2)
        image, boxes = self.get_data(self.anno_line[index], self.resize_hw, self.data_augmen)

        image = self.preprocess_img(image)  # 对boxes和image进行归一化操作,归一化到【0,1】
        if len(boxes) != 0:
            boxes[:, [0, 2]] = boxes[:, [0, 2]] / self.resize_hw[1]
            boxes[:, [1, 3]] = boxes[:, [1, 3]] / self.resize_hw[0]

        # 3.正常来说，在读取数据的时候只获取image和box就结束了。但这里为了求loss的时候方便，所以多求了一个y_true
        y_true = self.get_target(boxes)
        return image, boxes, y_true

    def get_data(self, anno_line, resize_hw, data_aug):
        """ line:读取txt文件每一行的信息 """
        line = anno_line.split(" ")
        img_path = line[0]

        image = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), 1)
        ih, iw = image.shape[:2]
        ratio_h, ratio_w = resize_hw[0] / ih, resize_hw[1] / iw

        # case1:数据只进行缩放，不进行其他操作
        if not data_aug:
            image = cv2.resize(image, (resize_hw[1], resize_hw[0]))
            boxes = [list(map(int, box.split(","))) for box in line[1:]]  # boxes是相对于原图而言的
            for box in boxes:  # 把boxes映射到resize之后的图
                box[0], box[2] = box[0] * ratio_w, box[2] * ratio_w
                box[1], box[3] = box[1] * ratio_h, box[3] * ratio_h
            return image, boxes
        # case2:进行数据增强操作
        else:
            image = cv2.resize(image, (resize_hw[1], resize_hw[0]))
            boxes = [list(map(int, box.split(","))) for box in line[1:]]
            for box in boxes:
                box[0], box[2] = box[0] * ratio_w, box[2] * ratio_w
                box[1], box[3] = box[1] * ratio_h, box[3] * ratio_h
            image, boxes = self.data_augment(image, boxes, h=resize_hw[0], w=resize_hw[1])
            return image, boxes

    def data_augment(self, image, boxes, h, w):
        random_flag = random.randint(0, 10)

        # case1:图像x轴对称（即：上下翻转）
        if random_flag == 2:
            image = cv2.flip(image, 0)
            for box in boxes:
                box[1], box[3] = h - box[3], h - box[1]
            boxes = np.array(boxes, dtype=np.float32)
            return image, boxes

        # case2: 图像y轴对称，（即：左右翻转）
        elif random_flag == 3:
            image = cv2.flip(image, 1)
            for box in boxes:
                box[0], box[2] = w - box[2], w - box[0]
            boxes = np.array(boxes, dtype=np.float32)
            return image, boxes

        # case3: 图像x,y轴对称
        elif random_flag == 4:
            image = cv2.flip(image, -1)
            for box in boxes:
                box[0], box[2] = w - box[2], w - box[0]
                box[1], box[3] = h - box[3], h - box[1]
            boxes = np.array(boxes, dtype=np.float32)
            return image, boxes

        # case4: 图像旋转180度
        elif random_flag == 5:
            image = cv2.rotate(image, cv2.ROTATE_180)
            for box in boxes:
                box[0], box[1], box[2], box[3] = w - box[2], h - box[3], w - box[0], h - box[1]
            boxes = np.array(boxes, dtype=np.float32)
            return image, boxes

        else:
            boxes = np.array(boxes, dtype=np.float32)
            return image, boxes

    def preprocess_img(self, image):
        image = image.astype("float32")
        image /= 255.
        image = np.transpose(image, (2, 0, 1))  # h,w,c--->c,h,w
        return image

    def get_near_points(self, x, y, i, j):
        sub_x = x - i
        sub_y = y - j
        if sub_x > 0.5 and sub_y > 0.5:
            return [[1, 1]]
        elif sub_x < 0.5 and sub_y > 0.5:
            return [[-1, 1]]
        elif sub_x < 0.5 and sub_y < 0.5:
            return [[-1, -1]]
        else:
            return [[1, -1]]

    def get_target(self, boxes):
        print("boxes:", boxes.shape)
        num_layers = len(self.anchors_mask)
        input_shape = np.array(self.resize_hw, dtype='int32')
        grid_shapes = [input_shape // {0: 32, 1: 16, 2: 8}[l] for l in range(num_layers)]

        # 1.y_true是一个列表【（3,16,16,5+nc）， （3,32,32,5+nc）， （3, 64,64, 5+nc）】的全零矩阵
        y_true = [np.zeros((3, grid_shapes[l][0], grid_shapes[l][1], self.bbox_attrs), dtype='float32')
                  for l in range(num_layers)]

        ## 2.box_best_ratio是框的调整系数
        # box_best_ratio = [np.zeros((3, grid_shapes[l][0], grid_shapes[l][1]), dtype='float32')
        #                   for l in range(num_layers)]

        if len(boxes) == 0:
            return y_true

        for l in range(num_layers):
            in_h, in_w = grid_shapes[l]
            print("ih, iw:", in_h, in_w)
            anchors = self.anchors / {0: 32, 1: 16, 2: 8}[l]  # anchors映射到feature map上
            print("anchors:", anchors.shape)
            batch_target = np.zeros_like(boxes)  # 创建一个全零矩阵、方便对中心、宽高进行调整,shape为[n_box, 5]

            batch_target[:, [0, 2]] = boxes[:, [0, 2]] * in_w  # 在读取数据的时候已经把boxes归一化到【0,1】之间
            batch_target[:, [1, 3]] = boxes[:, [1, 3]] * in_h  # 所以在这里要乘以feature map的宽高，将中心坐标对应到每一层feature map上
            batch_target[:, 4] = boxes[:, 4]

            print("batch_target:", batch_target.shape)
            ratios_of_gt_anchors = np.expand_dims(batch_target[:, 2:4], axis=1) / np.expand_dims(anchors, axis=0)
            ratios_of_anchors_gt = np.expand_dims(anchors, 0) / np.expand_dims(batch_target[:, 2:4], 1)
            ratios = np.concatenate([ratios_of_gt_anchors, ratios_of_anchors_gt], axis=-1)

            max_ratios = np.max(ratios, axis=-1)
            print("max_ratios:", max_ratios.shape)
            for t, ratio in enumerate(max_ratios):
                over_threshold = ratio < 4  # over_threshold：[1,9]  bool值
                over_threshold[np.argmin(ratio)] = True  # 为了防止没有一个值为True，所以设置最小的那个为true
                for k, mask in enumerate(self.anchors_mask[l]):
                    if not over_threshold[mask]:
                        continue

                    i = int(np.floor(batch_target[t, 0]))
                    j = int(np.floor(batch_target[t, 1]))
                    offsets = self.get_near_points(x=batch_target[t, 0], y=batch_target[t, 1], i=i, j=j)
                    for offset in offsets:
                        local_i = i + offset[0]
                        local_j = j + offset[1]

                        # 增加越界约束，超过边界之后就跳出循环
                        if local_i >= in_w or local_i < 0 or local_j >= in_h or local_j < 0:
                            continue

                        

                        c = int(batch_target[t, 4])  # 取出真实框的种类
                        print("class:", c)
                        y_true[l][k, local_j, local_i, 0] = batch_target[t, 0]
                        y_true[l][k, local_j, local_i, 1] = batch_target[t, 1]
                        y_true[l][k, local_j, local_i, 2] = batch_target[t, 2]
                        y_true[l][k, local_j, local_i, 3] = batch_target[t, 3]
                        y_true[l][k, local_j, local_i, 4] = 1
                        y_true[l][k, local_j, local_i, c + 5] = 1                # 第一个类别的位置为0+5，第一个类别的位置为1+5
                        # box_best_ratio[l][k, local_j, local_i] = ratio[mask]

        return y_true


if __name__ == '__main__':
    with open(r"E:\私人文件\V3软件标注\定位\tem\train.txt", "r") as f:
        lines = f.readlines()
    anchors = np.array(
        [[10., 13.], [16., 30.], [33., 23.], [30., 61.], [62., 45.], [59., 119.], [116., 90.], [156., 198.],
         [373., 326.]])
    data = GetDatasets(lines, [512, 1024], 3, anchors, True)
    data.__getitem__(2)