以上介绍了使用${\rm PyTorch}$加载数据时的整体结构，下面就${\rm VOC}$标注格式的数据集和${\rm COCO}$标注格式的数据集分别介绍相应的处理流程。

数据集的目录结构如下图。在制作自己的数据集时，首先我们要依照下图格式建立相应的文件夹。然后如上图分别执行对应的四个步骤。其中在画完框后会自动弹出来一个对话框，此时我们需要输入该标注目标的类别。最后根据标注信息将自动生成.xml文件。

其中，第一个文件夹用于存放数据集的标注信息，以.xml文件保存。我们以目标检测部分的标注信息介绍文件内的具体内容，以000002.xml为例。

第二个文件夹内的Main文件夹存放着目标检测相关的文件。其中train.txt、val.txt、test.txt和trainval.txt分别存放了训练集、验证集、测试集和训练验证集的图像名称。第三个文件夹内存放图像本身，如上述标注文件对应的图像000002.jpg为下图。

第四个文件夹和第五个文件夹内存放的图像分割的相关标注信息。首先我们来定义文件解析类，其参数是ElementTree类的对象（用于解析${\rm xml}$文件的类），返回是对应文件所包含的标注信息。

class VOCAnnotationTransform:def __init__(self):# 将类别标签转换为对应的数字标签self.class_to_ind = dict(zip(VOC_CLASSES, range(len(VOC_CLASSES))))def __call__(self, target, width, height):res = []for obj in target.iter('object'):# 目标类别name = obj.find('name').text.lower().strip()# 标注框bbox = obj.find('bndbox')pts = ['xmin', 'ymin', 'xmax', 'ymax']bndbox = []for i, pt in enumerate(pts):# 将坐标值缩放到[0,1]内cur_pt = int(bbox.find(pt).text) - 1cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / heightbndbox.append(cur_pt)# 获取类别和标注框信息并添加到结果label_idx = self.class_to_ind[name]bndbox.append(label_idx)res += [bndbox]return res  # res=[[xmin, ymin, xmax, ymax, label_ind], ... ]

然后定义数据加载类。在数据加载时，只有当使用到该条数据时我们才将其加载到内存，在函数pull_item函数内实现。最后通过__getitem__函数返回指定index的数据。

class VOCDetection(data.Dataset):def __init__(self, root, target_transform=VOCAnnotationTransform()):# 数据集根目录self.root = root# 调用解析类self.target_transform = target_transform# 文件路径self.annopath = osp.join('%s', 'Annotations', '%s.xml')self.imgpath = osp.join('%s', 'JPEGImages', '%s.jpg')# self.ids=((数据集根目录, 文件名),...)，作用是与上面“文件路径”变量组合称为完整路径self.ids = list()for line in open(osp.join(self.root, 'ImageSets', 'Main', "trainval" + '.txt')):self.ids.append((self.root, line.strip()))def __getitem__(self, index):im, gt, h, w = self.pull_item(index)# 返回数据return im, gtdef __len__(self):return len(self.ids)def pull_item(self, index):# 根据index取出某一条(数据集根目录, 文件名)img_id = self.ids[index]# 组合成完成路径后解析xml文件和读取图像target = ET.parse(self.annopath % img_id).getroot()img = cv2.imread(self.imgpath % img_id)height, width, channels = img.shape# xml解析if self.target_transform is not None:target = self.target_transform(target, width, height)return torch.from_numpy(img).permute(2, 0, 1), target, height, width

在使用各${\rm API}$前，我们需要实例化COCO类，它接受的参数为标注文件的路径，返回类的对象。以本文的内容为例，首先我们使用以下语句初始化COCO类的对象。这里使用的是${\rm COCO\ 2017}$数据集中对应的训练集部分。

coco = COCO(os.path.join(root, 'annotations', 'instances_{}.json'.format('train2017')))

然后，我们就可以通过COCO类的对象调用各种${\rm API}$函数。其中，在本文将会使用的${\rm API}$函数包括：

1. coco.imgToAnns将图像的索引与其标注信息相关联，执行后的效果是给定指定的图像索引可以返回该图像对应的所有标注信息，coco.imgToAnns.keys()返回所有的图像的索引（给数据集中的每幅图像赋值一个索引，用于后续与其标注和类别信息相关联），然后再通过指定图像的索引就可以访问其相关的标注信息。如下图是各索引之间的相互关联。
   
   注意每个annotation里面仅对应于一个目标的标注信息。如果我们使用coco.imgToAnns.keys()[index]，则 将区域索引为${\rm index}$的图像对应的标注信息的索引，是一个列表。然后根据每个标注信息的索引去寻找每一个目标的标注信息。如上图，bbox即为本文中我们所需要的目标检测的标注信息。
2. coco.getAnnIds(imgIds=imgIds)就是根据参数值取指定索引图像的标注信息，返回一个列表。
3. 得到标注信息的索引的列表后，我们就可以使用coco.loadAnns(ids=ann_ids)返回指定标注索引的标注内容，其中同时包括目标检测和图像分割的内容。也就是上图中的annotations{}部分。
4. 同时，根据coco.loadImgs(ids=img_ids)就可以获得指定图像索引的图像信息，其格式如下：

则现在我们可以完成整个数据加载类的书写。首先我们需要注意的是，由于${\rm COCO}$数据集中的索引并不是连续的，如图：最左边是原始的索引，中间是经过处理后的索引，最后一列表示具体的类别。首先，我们需要根据该文件的内容使用中间一列的索引作为最后的索引。定义如下函数：

def get_label_map(label_file):label_map = {}labels = open(label_file, 'r')for line in labels:ids = line.split(',')# 返回字典形式，如上图中的为{...,11:11,13:12,14:16,...}label_map[int(ids[0])] = int(ids[1])return label_map

和${\rm VOC}$数据集的加载流程一致，首先我们定义解析类COCOAnnotationTransform，传入参数是标注信息，即上文提到的一系列的annotations{}。然后返回形式同VOCAnnotationTransform类一致。

class COCOAnnotationTransform:def __init__(self):self.label_map = get_label_map(osp.join('data', 'coco_labels.txt'))def __call__(self, target, width, height):scale = np.array([width, height, width, height])res = []for obj in target:if 'bbox' in obj:# 将(x,y,w,h)->(xmin,ymin,xmax,ymax)bbox = obj['bbox']bbox[2] += bbox[0]bbox[3] += bbox[1]# 将坐标值缩放到[0,1]内final_box = list(np.array(bbox)/scale)# 获取类别和标注框信息并添加到结果label_idx = self.label_map[obj['category_id']] - 1final_box.append(label_idx)res += [final_box]else:print("no bbox problem!")return res  # [[xmin, ymin, xmax, ymax, label_idx], ... ]

然后定义数据加载类，相应内容同上。

class COCODetection(data.Dataset):def __init__(self, root, target_transform=COCOAnnotationTransform()):# 图像数据集根目录self.root = osp.join(root, 'train2017')# 参数为标注文件路径，返回COCO类的对象self.coco = COCO(osp.join(root, 'annotations', 'instances_{}.json'.format('train2017')))# 获取每张图片的索引，同时将图片索引与标注信息相关联self.ids = list(self.coco.imgToAnns.keys())# 调用解析类self.target_transform = target_transformdef __getitem__(self, index):im, gt, h, w = self.pull_item(index)# 返回数据return im, gtdef __len__(self):return len(self.ids)def pull_item(self, index):# 获得图像索引img_id = self.ids[index]# 获得指定图像索引的标注信息，返回一个列表，列表的每一个值表示一个目标的标注信息ann_ids = self.coco.getAnnIds(imgIds=img_id)# 根据标注信息的索引返回其具体的标注内容target = self.coco.loadAnns(ann_ids)# 获得图像的完整路径path = osp.join(self.root, self.coco.loadImgs(img_id)[0]['file_name'])assert osp.exists(path), 'Image path does not exist: {}'.format(path)img = cv2.imread(path)# 解析标注信息height, width, _ = img.shapeif self.target_transform is not None:target = self.target_transform(target, width, height)return torch.from_numpy(img).permute(2, 0, 1), target, height, width