自定义数据集思路（以PASCAL VOC数据集为例）

PASCAL VOC数据集文件以及读取方法详细介绍

一直在使用PASCAL VOC在做目标检测和分割任务，但是关注点都在网络构造和评价指标上面，数据集方面就直接按照别人的代码使用了。
最近准备开始使用COCO了，但是开始了解COCO的时候发现数据集文件内容还是不太了解，于是重新去学习了一下PASCAL VOC的文件和读取。

首先，主要参考的是CSDN博主 太阳花的小绿豆
的博文。该数据集主要用于图像分类(Object Classification)，目标检测(Object Detection)，目标分割(Object Segmentation)，行为识别(Action Classification) 等。
该数据集的其他信息可以去官网查看，这里不做介绍。

这里使用的PASCAL VOC数据集的版本主要是training/validation data (2GB tar file)

该数据集文件的目录：

VOCdevkit
    └── VOC2012
         ├── Annotations               所有的图像标注信息(XML文件)
         ├── ImageSets    
         │   ├── Action                人的行为动作图像信息
         │   ├── Layout                人的各个部位图像信息
         │   │
         │   ├── Main                  目标检测分类图像信息
         │   │     ├── train.txt       训练集(5717)
         │   │     ├── val.txt         验证集(5823)
         │   │     └── trainval.txt    训练集+验证集(11540)
         │   │
         │   └── Segmentation          目标分割图像信息
         │         ├── train.txt       训练集(1464)
         │         ├── val.txt         验证集(1449)
         │         └── trainval.txt    训练集+验证集(2913)
         │ 
         ├── JPEGImages                所有图像文件
         ├── SegmentationClass         语义分割png图（基于类别）
         └── SegmentationObject        实例分割png图（基于目标）

Annotations 文件就是该数据集所有图像的标注信息，包括了一张图片的文件名，图片中的目标及名称等，该目标的bbox位置等,共17125条。
JPEGImages文件夹里面含有该数据集的所有图片，共17125张。所以每张图片都对应一条annotation。

ImageSets文件夹就是一种功能性的文件夹，里面根据VOC数据集所能实现的任务对图片进行了分类。比如Action为行为检测，Layout为人体部位检测，Main为目标检测，Segmentation为分割任务。注意，Segmentation里面索引对应的图片既可以做基于类别的分割也可以做基于目标的分割。

由于做分割任务时，分割的标注信息无法储存在annotation里面，所以需要单独标注并储存，也就是最后的SegmentationClass和SegmentationObject文件内的内容。

上面就是数据集文件的结构和内容。

实际上，在数据集文件目录中，不包含真正图片和标注的文件夹只有ImageSets，可以发现里面的文件全是txt文件，文本内容就是部分图片的索引。
所以在代码中使用该数据集的逻辑就是，根据要求在ImageSets文件先选择对应功能的文件夹，里面有train.txt和val.txt。如果是训练，就在train.txt里面抽取一条记录，该记录就是对应图片的索引（实际上就是图片的名称）。然后根据该索引Annotations中找到对应annotation，解析该annotations获得对应图片的名称，如果是检测任务，则再获取相应目标的bbox信息。然后根据图片名称再JPEGImages里面找到原图片。如果是分割任务，则还需去SegmentationClass或者SegmentationObject文件中找到对应的标注图片。这样，一张图片和对应的标注图片（信息）就都获得了，数据集类的__getitems__方法就实现了。

以PASCAL VOC数据集为例的自定义数据集思路

首先，自定义数据集必须定义的三个函数：__init__()、__getitem__()和__len__()。（其实len()不必要）
每个函数所需要发挥的作用如下：

__init__() : 为__getitem__()方法铺路。根据给出的数据集root路径获得原图片路径，target路径和train文本或者val文本路径。类似于train文本的格式，将训练集每一张图片的路径储存到一个列表中，类似的，将target信息一并储存到另一个列表中。最终效果就是，给定一个索引，就能通过图片路径列表获得图片，通过target列表获得对应的target信息。

因为是PASCAL VOC数据集，如果想要实现一个普适性的数据集，就需要将语义分割和目标检测的target合并在一个target内，就类似于这里实现的方法。

__getitem__() : 通过索引获得数据集中的图片和对应的target。这个方法的实现全在__init__()方法的基础上，也就是如果__init__()铺路铺的好，此方法的实现就十分简单。

__len__() : 获取数据集的大小。这个函数可以不实现，实现起来也非常简单，一行代码就可以实现。

所以最重要的函数其实是__init__()函数。其中不仅包括之前说的路径处理，还包括target信息的获获取和处理

__init__():

    def __init__(self,root,year='2012',txt_name='train',transforms=None):
        # 对于根目录，先检测存在
        root = os.path.join(root,"VOCdevkit")
        assert os.path.exists(root), f"dataset does not exist the directory {root}"

        # 检测数据集版本和必要内容是否存在
        root = os.path.join(root,f"VOC{year}")
        assert os.path.exists(root), f"the vision of dataset {year} does not exist"

        # 数据集文件存在后，生成所需文件路径,并检测
        image_dir = os.path.join(root,"JPEGImages")
        xml_dir = os.path.join(root,"Annotations")
        mask_dir = os.path.join(root,"SegmentationObject")

        assert os.path.exists(image_dir), "the needed images file does not exist"
        assert os.path.exists(xml_dir), "the needed annotation file does not exist"
        assert os.path.exists(mask_dir), "the neeses mask file does not exist"

        txt_path = os.path.join(root, "ImageSets", "Segmentation", txt_name+".txt")
        assert os.path.exists(txt_path), "the text file does not exsit"

        # 读取text的内容，也就是训练或者验证图片的名称
        with open(txt_path,"r") as f:
            file_names = [x.strip() for x in f.readlines() if len(x.strip()) > 0]

        self.images_path = [os.path.join(image_dir, x+".jpg") for x in file_names]
        self.xmls_path = [os.path.join(xml_dir,x+".xml") for x in file_names]
        self.masks_path = [os.path.join(mask_dir,x+".png") for x in file_names]

        # 读取类别的indices文件
        json_file = 'pascal_voc_indices.json'
        assert os.path.exists(json_file), "{} file not exist.".format(json_file)
        with open(json_file, 'r') as f:
            idx2classes = json.load(f)
            self.class_dict = dict([(v, k) for k, v in idx2classes.items()])


        self.xmls_info = []
        self.objects_bboxes = []
        self.mask = []

        # 将一张图片的image、xml和mask打包在一起
        for index, (img_path, xml_path, mask_path) in enumerate(zip(self.images_path,self.xmls_path,self.masks_path)):
            assert os.path.exists(img_path), f"not find {img_path}"
            assert os.path.exists(xml_path), f"not find {xml_path}"
            assert os.path.exists(mask_path), f"not find {mask_path}"

            # 读取xml中的bbox信息
            with open(xml_path) as fid:
                xml_str = fid.read()
            xml = etree.fromstring(xml_str)
            obs_dict = parse_xml(xml)["annotation"] # 解析xml文件，变成一个字典
            obs_bboxes = parse_object(obs_dict, xml_path, self.class_dict, index) # 解析字典中的object节点，获取其中的必要信息。

            # 读取SegmentationObject,并处理mask中的白边
            instances_mask = Image.open(mask_path)
            instances_mask = np.array(instances_mask)
            instances_mask[instances_mask == 255] = 0

            self.xmls_info.append(obs_dict)
            self.objects_bboxes.append(obs_bboxes)
            self.mask.append(instances_mask)

        self.transforms = transforms

parse_xml()用于解析xml对象，他将返回一个字典，该字典是按照节点（树）的形式组织的。该函数模板是通用的。

def parse_xml(xml):
    if len(xml) == 0:
        return {xml.tag: xml.text}

    result = {}
    for child in xml:
        child_result = parse_xml(child)
        if child.tag != "object":
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}

parse_object()函数如下，该函数模板是通用的，可以根据需要修改返回的内容。

def parse_object(data, xml_path, class_dict, idx):
    boxes = []
    labels = []
    difficult = []

    assert "object" in data, f"{xml_path} lack of object information"
    for obj in data["object"]:
        xmin = float(obj["bndbox"]["xmin"])
        xmax = float(obj["bndbox"]["xmax"])
        ymin = float(obj["bndbox"]["ymin"])
        ymax = float(obj["bndbox"]["ymax"])


        if xmax <= xmin or ymax <= ymin:
            print(f"warning:in {xml_path} xml, there are some bbox w/h <= 0")
            continue

        boxes.append([xmin, ymin, xmax, ymax])
        labels.append(int(class_dict[obj["name"]]))
        if "difficult" in obj:
            difficult.append(int(obj["difficult"]))
        else:
            difficult.append(0)

    boxes = torch.as_tensor(boxes, dtype=torch.float32)
    labels = torch.as_tensor(labels, dtype=torch.int64)
    difficult = torch.as_tensor(difficult, dtype=torch.int64)

    image_id = torch.as_tensor([idx])
    area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

    return {"boxes": boxes,
            "labels": labels,
            "difficult": difficult,
            "image_id": image_id,
            "area": area}

所以至此，我们的铺路工作就全部做完了，接下来就是实现简单的__getitem__()和__len__():

    def __getitem__(self, idx):

        img = Image.open(self.images_path[idx]).convert('RGB')
        target, _, _  = self.get_annotations(idx)  # 将分割的target和目标检测的target组合在一起。

        if self.transforms is not None:
            img,target = self.transforms(img,target)

        return img,target

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

get_annotation()方法就是将mask信息和bboxes信息组织在一起。
parse_mask()函数将实例分割target图像中的每一个目标单独提取出来，方便做分割任务。

    def get_annotations(self,idx):
        h, w = self.get_height_and_width(idx)
        target = self.objects_bboxes[idx]
        masks = self.parse_mask(idx)
        target["masks"] = masks
        return target, h, w
    def get_height_and_width(self, idx):
        data = self.xmls_info[idx]
        h = int(data["size"]["height"])
        w = int(data["size"]["width"])
        return h, w
    
    def parse_mask(self, idx: int):
        mask = self.masks[idx]
        c = mask.max()  # 有几个目标最大索引就等于几
        masks = []
        # 对每个目标的mask单独使用一个channel存放
        for i in range(1, c+1):
            masks.append(mask == i)
        masks = np.stack(masks, axis=0)
        return torch.as_tensor(masks, dtype=torch.uint8)

至此，PASCAL VOC自定义数据集就定义完成了。可以用下面语句测试一下：

if __name__ == "__main__":
    dataset = VOCInstances(root=r"./data")
    print(len(dataset))
    # 调用__getitem__方法，获取idx为0的图片和target
    d1 = dataset[0]
    # 展示图片
    d1[0].show()
    # 展示target字典
    print(d1[1])

总结

自定义数据集前要清楚所要实现的三个函数分别需要执行什么功能，需要处理什么数据。
init函数主要处理好路径并对target做好解析。getitem函数根据init函数的处理结果进行简单的返回。len函数最简单，只需要返回数据集大小。

以上内容主要参考博主 太阳花的小绿豆
的博文，最初代码可以看该博主的Github里的代码，以上代码为本人修改后的代码。