划分训练集和测试集

annotation_path = r'2007_train.txt'
    with open(annotation_path, encoding='utf-8') as f:
        lines = f.readlines()
    np.random.seed(10101)
    np.random.shuffle(lines)
    np.random.seed(None)
    val = 0.1
    num_val = int(len(lines)*val)
    num_train = len(lines) - num_val

定义自己的数据集处理代码，数据处理主要继承torch.utils.data中的Dataset函数【注意：这里只是预处理数据集，还不是加载数据集，加载预处理后的数据集需要用torch.utils.data中的DataLoader进行加载】

class MyDatasets(Dataset):
    def __init__(self, train_line, image_size, is_train):
        super(MyDatasets, self).__init__()
        self.train_line = train_line  # 训练集的长度
        self.train_batches = len(train_line)  # 总batch数量
        self.image_size = image_size  # 图像大小
        self.is_train = is_train  # 是否训练

继承Dataset函数

继承Dataset父类函数主要修改两个成员函数，def __len__(self):【用来返回数据集长度】和def __getitem__(self, index):【用来返回数据集和标签】。注意后者是一个可迭代的【看形参index就可知道】。

    def __len__(self):  # 返回数据集的长度
        return self.train_batches
    def __getitem__(self, index):  # 返回数据集和标签
        lines = self.train_line
        if self.is_train:
            img, y = self.get_data(lines[index], self.image_size[0:2], random=False)

我用自己的数集来演示一下上述两个函数的输出结果【前提是已经划分好数据集了，将标签文件转为txt文件】

self.train_batches输出结果

In [1]: self.train_batches

Out[1]: 799

self.train_line输出结果

In [2]: lines = self.train_line

Out[3]:

['E:\\VOCdevkit/VOC2007/JPEGImages/23.jpg 3

7,27,134,209,0 233,1,387,210,0 20,220,167,414,0 258,226,348,413,0\n',

'E:\\VOCdevkit/VOC2007/JPEGImages/263.jpg

110,83,401,407,0\n',

'E:\\VOCdevkit/VOC2007/JPEGImages/117.jpg

205,115,342,329,0\n',

···········································

可以看出self.train_line是将txt文件中的数据进行进行了读取VOCdevkit/VOC2007/JPEGImages为路径23.jpg图片名称，37,27,134,209,0 233,1,387,210,0 20,220,167,414【bbox信息】,0【类，该图像中的类别】，代码中的lines类型是列表

定义数据集处理函数：

def get_data(self, annotation_line, input_shape)：
    line = annotation_line.split()
    image = Image.open(line[0])  # line[0]是图片路径，line[1:]是框和标签信息
    iw, ih = image.size  # 真实输入图像大小
    h, w = input_shape  # 网络输入大小
    box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])  # 将box信息转为数组

将上述lines【txt文件中数据集】放入get_data函数中，即对应的形参annotation_line。这里我只加从lines列表中载一张图像进行演示。

annotation_line.split()是将图像路径、图像名称、bbox和类别进行划分

所以line = annotation_line.split()的输出为：

In [5]: line

Out[5]:

['E:\\VOCdevkit/VOC2007/JPEGImages/23.jpg',

'37,27,134,209,0',

'233,1,387,210,0',

'20,220,167,414,0',

'258,226,348,413,0']

打开图片

image = Image.open(line[0])  # line[0]是图片路径，line[1:]是框和标签信息

可以对图片进行显示一下给大家看看

image.show()

读取图像和bbox信息

iw, ih = image.size  # 真实输入图像大小
h, w = input_shape  # 网络输入大小
box = np.array([np.array(list(map(int, box.split(',')))) for box in line[1:]])  # 将box信息转为数组

iw,ih指的是输入图像的真实大小

h,w指的是网络输入大小，比如SSD是300*300，yolo是416*416，这里我以网络输入300*300为例，即h = w = 300

In [9]: iw,ih=image.size

In [10]: iw,ih

Out[10]: (416, 416)

------------------------------------------------------------------

h=w=300

box输出信息 ：注意看下面的输出，这个数组或者说矩阵中，一共5列，从第一列到第四列是bbox信息，最后一列是类标签。并且在前四列中的bbox信息中，1和2列是bbox左上角坐标，3和4列是右下角坐标。这里非常重要，后面还要用到。

In [13]: box

Out[13]:

array([[ 37,  27, 134, 209,   0],

      [233,   1, 387, 210,   0],

      [ 20, 220, 167, 414,   0],

      [258, 226, 348, 413,   0]])

因为我们的网络输入是300*300，但这里的图像实际大小为416*416，那么我们需要对原始图像进行缩放，使其适合我们的网络输入。

图像缩放

先计算一下缩放比例，也就是原始图像相对于输入图像大小，缩放后的图像不产生形变：

scale = min(w / iw, h / ih)
保证长或宽，符合目标图像的尺寸
nw = int(iw * scale)
nh = int(ih * scale)
dx = (w - nw) // 2  
dy = (h - nh) // 2

In [15]: scale

Out[15]: 0.7211538461538461

In [18]: nw,nh

Out[18]: (300, 300)

image = image.resize((nw, nh), Image.BICUBIC) # 采用双三次插值算法缩小图像

创建一个新的灰度图像，缩放后的图像可能不能满足网络大小，所以可以给周边补充一些灰度条。

new_image = Image.new('RGB', (w, h), (128, 128, 128))

可以看一下这个新创造的图，new_image.show()

再将灰度条和裁剪后图合在一起

1. new_image.paste(image, (dx, dy))
2. image_data = np.array(new_image, np.float32)