SSD源码解读——数据读取
之前，对SSD的论文进行了解读，可以回顾之前的博客：https://www.cnblogs.com/dengshunge/p/11665929.html。
为了加深对SSD的理解，因此对SSD的源码进行了复现，主要参考的github项目是ssd.pytorch。同时，我自己对该项目增加了大量注释：https://github.com/Dengshunge/mySSD_pytorch
搭建SSD的项目，可以分成以下三个部分：
接下来，本篇博客重点分析数据读取。
一、整体框架
SSD的数据读取环节，同样适用于大部分目标检测的环节，具有通用性。为了方便理解，本项目以VOC2007+2012为例。因此，数据读取环节，通常是按照以下步骤展开进行：
- 函数入口；
- 图片的读取和xml文件的读取；
- 对GT框进行处理；
- 数据增强；
- 辅助函数。
二、具体实现细节
2.1 函数入口
数据读取的函数入口在train.py文件中：
if args.dataset == 'VOC': train_dataset = VOCDetection(root=args.dataset_root) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=args.batch_size, num_workers=4, collate_fn=detection_collate, shuffle=True, pin_memory=True)
可以看到，首先通过函数 VOCDetection() 来对VOC数据集进行初始化，再使用函数 DataLoader() 来实现对数据集的读取。这一步与常见的分类网络形式相同，但不同的是，多了collate_fn这一参数，后续会对此进行说明。
2.2 图片与xml文件读取
首先，我们先看看函数VOCDetection() 的初始化函数__init__()。在__init__中包含了需要传入的几个参数，image_sets（表示VOC使用到的数据集），transform（数据增强的方式），target_transform（GT框的处理方式）。
class VOCDetection(): """VOC Detection Dataset Object input is image, target is annotation Arguments: root (string): filepath to VOCdevkit folder. image_set (string): imageset to use (eg. 'train', 'val', 'test') transform (callable, optional): transformation to perform on the input image 图片预处理的方式，这里使用了大量数据增强的方式 target_transform (callable, optional): transformation to perform on the target `annotation` (eg: take in caption string, return tensor of word indices) 真实框预处理的方式 """ def __init__(self, root, image_sets=[('2007', 'trainval'), ('2012', 'trainval')], transform=SSDAugmentation(size=config.voc['min_dim'], mean=config.MEANS), target_transform=VOCAnnotationTransform()): self.root = root self.image_set = image_sets self.transform = transform self.target_transform = target_transform self._annopath = os.path.join('%s', 'Annotations', '%s.xml') self._imgpath = os.path.join('%s', 'JPEGImages', '%s.jpg') self.ids = [] # 使用VOC2007和VOC2012的train作为训练集 for (year, name) in self.image_set: rootpath = os.path.join(self.root, 'VOC' + year) for line in open(os.path.join(rootpath, 'ImageSets', 'Main', name + '.txt')): self.ids.append([rootpath, line[:-1]])
首先，为什么image_sets是这样的形式呢？因为VOC具有固定的文件夹路径，利用这个参数和配合路径读取，可以读取到txt文件，该txt文件用于制定哪些图片用于训练。此外，还需要设置参数self.ids，这个list用于存储文件的路径，由两列组成，"VOC/2007"和图片名称。通过这两个参数，后续可以配合函数_annopath()和_imgpath()可以读取到对应图片的路径和xml文件。
在pytorch中，还需要相应的函数来对读取图片与返回结果，如下所示。其中，重点是pull_iterm函数。
def __getitem__(self, index): im, gt = self.pull_item(index) return im, gt def __len__(self): return len(self.ids) def pull_item(self, index): img_id = tuple(self.ids[index]) # img_id里面有2个值 target = ET.parse(self._annopath % img_id).getroot() # 获得xml的内容，但这个是具有特殊格式的 img = cv2.imread(self._imgpath % img_id) height, width, _ = img.shape if self.target_transform is not None: # 真实框处理 target = self.target_transform(target, width, height) if self.transform is not None: # 图像预处理，进行数据增强,只在训练进行数据增强,测试的时候不需要 target = np.array(target) img, boxes, labels = self.transform(img, target[:, :4], target[:, 4]) # 转换格式 img = img[:, :, (2, 1, 0)] # to rbg target = np.hstack((boxes, np.expand_dims(labels, axis=1))) return torch.from_numpy(img).permute(2, 0, 1), target
该函数pull_item()，首先读取图片和相应的xml文件；接着对使用类VOCAnnotationTransform来对GT框进行处理，即读取GT框坐标与将坐标归一化；然后通过函数SSDAugmentation()对图片进行数据增强；最后对对图片进行常规处理（交换通道等），返回图片与存有GT框的list。
2.3 对GT框进行处理
接着，需要讲一讲这个类VOCAnnotationTransform的作用，其定义如下。self.class_to_ind是一个map，其key是类别名称，value是编号，这个对象的作用是，读取xml时，能将对应的类别名称转换成label；在__call__()函数中，主要是xml读取的一些方式，值得一提的是，GT框的最表转换成了[0,1]之间，当图片尺寸变化了，GT框的坐标也能进行相应的变换。最后，res的每行由5个元素组成，分别是[x_min,y_min,x_max,y_max,label]。
class VOCAnnotationTransform(): ''' 获取xml里面的坐标值和label，并将坐标值转换成0到1 ''' def __init__(self, class_to_ind=None, keep_difficult=False): # 将类别名字转换成数字label self.class_to_ind = class_to_ind or dict(zip(VOC_CLASSES, range(len(VOC_CLASSES)))) # 在xml里面，有个difficult的参数，这个表示特别难识别的目标，一般是小目标或者遮挡严重的目标 # 因此，可以通过这个参数，忽略这些目标 self.keep_difficult = keep_difficult def __call__(self, target, width, height): ''' 将一张图里面包含若干个目标，获取这些目标的坐标值，并转换成0到1，并得到其label :param target: xml格式 :return: 返回List,每个目标对应一行,每行包括5个参数[xmin, ymin, xmax, ymax, label_ind] ''' res = [] for obj in target.iter('object'): difficult = int(obj.find('difficult').text) == 1 # 判断该目标是否为难例 # 判断是否跳过难例 if not self.keep_difficult and difficult: continue name = obj.find('name').text.lower().strip() # text是获得目标的名称,lower将字符转换成小写,strip去除前后空格 bbox = obj.find('bndbox') # 获得真实框坐标 pts = ['xmin', 'ymin', 'xmax', 'ymax'] bndbox = [] for i, pt in enumerate(pts): cur_pt = int(bbox.find(pt).text) - 1 # 获得坐标值 # 将坐标转换成[0,1]，这样图片尺寸发生变化的时候，真实框也随之变化，即平移不变形 cur_pt = cur_pt / width if i % 2 == 0 else cur_pt / height bndbox.append(cur_pt) label_idx = self.class_to_ind[name] # 获得名字对应的label bndbox.append(label_idx) res.append(bndbox) # [xmin, ymin, xmax, ymax, label_ind] return res # [[xmin, ymin, xmax, ymax, label_ind], ... ]
2.4 数据增强
还有一个重要的函数，即函数SSDAugmentation()，该函数的作用是作数据增强。论文中也提及了，数据增强对最终的结果提升有着重大作用。博客1和博客2具体讲述了数据增强的源码，讲得十分详细。在本项目中，SSDAugmentation()函数在data/augmentations.py中，如下所示。由于opencv读取读片的时候，取值范围是[0,255]，是int类型，需要将其转换为float类型，计算其GT框的正式坐标。然后对图片进行光度变形，包含改变对比度，改变饱和度，改变色调、改变亮度和增加噪声等。接着有对图片进行扩张和裁剪等。在此操作中，会涉及到GT框坐标的变换。最后，当上述变化处理完后，再对GT框坐标归一化，和resize图片，减去均值等。具体细节，可以参考两篇博客进行解读。
class SSDAugmentation(object): def __init__(self, size=300, mean=(104, 117, 123)): self.mean = mean self.size = size self.augment = Compose([ ConvertFromInts(), # 将图片从int转换成float ToAbsoluteCoords(), # 计算真实的锚点框坐标 PhotometricDistort(), # 光度变形 Expand(self.mean), # 随机扩张图片 RandomSampleCrop(), # 随机裁剪 RandomMirror(), # 随机镜像 ToPercentCoords(), Resize(self.size), SubtractMeans(self.mean) ]) def __call__(self, img, boxes, labels): return self.augment(img, boxes, labels)
2.5 辅助函数
在一个batch中，每张图片的GT框数量是不等的，因此，需要定义一个函数来处理这种情况。函数detection_collate()就是用于处理这种情况，使得一张图片能对应一个list，这里list里面有所有GT框的信息组成。
def detection_collate(batch): """Custom collate fn for dealing with batches of images that have a different number of associated object annotations (bounding boxes). 自定义处理在同一个batch,含有不同数量的目标框的情况 Arguments: batch: (tuple) A tuple of tensor images and lists of annotations Return: A tuple containing: 1) (tensor) batch of images stacked on their 0 dim 2) (list of tensors) annotations for a given image are stacked on 0 dim """ targets = [] imgs = [] for sample in batch: imgs.append(sample[0]) targets.append(torch.FloatTensor(sample[1])) return torch.stack(imgs, 0), targets
至此，已经将SSD的数据读取部分分析完。
