自定义目标检测数据集

前言

目标检测的数据集跟分类数据集不同，不仅仅包含图片和图片的类别信息，还要包括bounding box等一些额外的信息。因此，使用之前建立目录结构 + datasets.ImageFolder()的方法是不行的。

总的来看，自定义目标检测数据集要围绕原始图像、标注文件、入选名单三个要素，完成图像标注和自定义dataset类两项工作。

三个要素

  1. 原始图像，统统存放在一个目录下
  2. 标注文件，一般为一个xml或json，文件名与原始图像一一对应，保存对应图像的标注相关信息
  3. 入选名单，一个txt文档，记录分到train_set和val_set的名单，每行一个去掉后缀的文件名

两项工作

图像标注

  1. 安装lableImg

    pip3 install labelImg #中间遇到报错，参考here解决。

  2. 新建一个文件夹，包含img目录，annotation目录，class.txt，分别用于存放待标注图像，标注结果，类别清单。

    labelImg ./img ./class.txt #即可开始标注

  3. 标注完成后写一个随机抽样脚本，生成train_set和val_set入选清单两个txt文件。
    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    import os
    import random

    def main():
        random.seed(0)  # 设置随机种子，保证随机结果可复现

        files_path = "./annotation"
        assert os.path.exists(files_path), "path: '{}' does not exist.".format(files_path)

        val_rate = 0.5

        files_name = sorted([file.split(".")[0] for file in os.listdir(files_path)])
        files_num = len(files_name)
        val_index = random.sample(range(0, files_num), k=int(files_num*val_rate))
        train_files = []
        val_files = []
        for index, file_name in enumerate(files_name):
            if index in val_index:
                val_files.append(file_name)
            else:
                train_files.append(file_name)

        try:
            train_f = open("train.txt", "x")
            eval_f = open("val.txt", "x")
            train_f.write("\n".join(train_files))
            eval_f.write("\n".join(val_files))
        except FileExistsError as e:
            print(e)
            exit(1)

    if __name__ == '__main__':
        main()

自定义dataset类

自定义一个dataset类需要先继承torch.utils.data.Dataset类，然后重写__init__方法，__getitem__方法，__len__方法

__init__方法

主要有以下几项工作：

  1. 根据train_set.txt或者val_set.txt的内容建立self.xml_list，一个存放所有xml文件路径的list，并检查这些xml文件是否确实存在
  2. 根据描述类别信息的json文件建立self.class_dict
  3. 根据传入参数设定预处理函数self.transforms的行为，包括RandomFlip等augmentation操作以及to_tensor，normalize等操作，在getitem方法中也就是具体内容前执行。

__getitem__方法

__getitem__方法目的是定义用[idx]索引本类示例时返回什么。在这里，是返回一个tuple，内容为(image_tensor, img_attr_dict)。整个函数的行为就是为了得到这两个元素。

  1. 解析self.xml_list[idx]对应的xml，得到一个dict

    借助lxml的etree工具

  2. 根据dict中的filename找到对应的图像文件，并转化tensor，得到image_tensor

    借助PIL库的Image.open()函数

  3. 根据dict中的其他信息找到boxes，label等标注信息，并转为tensor，最后组装成img_attr_dict。

    多目标情形时，box作为key，对应一个list of tensor，每个tensor对应一个box，label等同理。一个典型的包含两个box的dict如下：
    {boxes:tensor(2x4),
    labels:tensor(2x1),
    image_id:tensor(1x1),
    area:tensor(2x1),
    iscrowd:tensor(2x1),
    }

__len__方法

返回self.xml_list的长度即可

one more thing

当torch.utils.data.DataLoader按照既定的batch_size去获取和组装图像时，不能用默认的collate_fn。

因为这个默认的collate_fn只是简单的调用torch.stack()将b个[c, h, w]的tensor组合成一个[b, c, h, w]的tensor。但是这里索引到的是b个(image_tensor, img_attr_dict)，直接stack将得到一个

[(image_tensor_0, img_attr_dict_0)
···
(image_tensor_b, img_attr_dict_b)]

这不是我们想要的。我们想要的是各个图像的像素值集中在一起，图像的标注信息集中在一起：

[(image_tensor_0,
…
image_tensor_b),

(img_attr_dict_0,
…,
img_attr_dict_b)]

所以要

1
2
def collate_fn(batch):
    return tuple(zip(*batch))