pytorch COCO2017 目标检测 DataLoader

Posted 2021-01-20 zi-wang

pytorch coco 目标检测 DataLoader实现

pytorch实现目标检测目标检测算法首先要实现数据的读入，即实现Dataset和DataLoader两个类。
借助pycocotools实现了CoCo2017用于目标检测数据的读取，并使用cv2显示。

分析

使用cv2显示读入数据，或者要送入到网络的数据应该有三个部分

1. 图像，Nx3xHeight x Width
2. BBs，NxMx4
3. 类型，NxMx1
   因此，可以将BBs和类型组成一个。Pytorch默认的数据类型是batchsize x nChanns x H x W。

在目标检测中，一般将图像进行缩放，使其尺寸满足一定要求，具体可以参考之前的博客。
也就是要实现一个Resizer()的类进行变换。此外，通常要对图像进行标准化处理，以及水平翻转等变换。因此，在实现Dataset时要
实现的变换有三个: Resizer()、Normilizer()和Augmenter()。

Python中图像数据读入一般都是 nChanns x H x W的numpy数组。常规的做法是使用Dataset中的transform对数据进行转换，
输出torch类型的数组。

由于CoCo数据集中图像的尺寸不一致，不能直接获得Nx3xHeight x Width类型的数组，因此要重写DataLoader中的collate_fn，
将一个minibatch中的图像尺寸调整一致。如果想要按照图像被缩放比例进行采样，就要重写DataLoader中的batch_sampler，
batch_sampler与DataLoader中的batch_size, shuffle, sampler, and drop_last参数是不兼容的，即在DataLoader中使用了batch_sampler
参数就不能再设置batch_size, shuffle, sampler, and drop_last参数。

从coco数据中读入图像、BBs以及类型

coco.getImgIds()返回了图像索引数组，可以分别结合coco.loadImgs()和coco.getAnnIds()分别获得图像、BBs和类型的具体信息。
要注意的事情有：

1. python中图像的读入的通常是numpy的uint8数组，需要转换成float类型，并除以255以使最大值为1.0；
2. coco数据中有80个类型，但是给的标签值最大为90，说明并不连续，需要设置新的标签，新的标签要从0到79，一定从0开始。
3. coco数据集中有些图片的BBs标签高宽小于1，标注的问题，要注意舍去

下面就是一个简单的SimpleCoCoDataset类

class SimpleCoCoDataset(Dataset):
    def __init__(self, rootdir, set_name=‘val2017‘, transform=None):
        self.rootdir, self.set_name = rootdir, set_name
        self.transform = transform
        self.coco = COCO(os.path.join(self.rootdir, ‘annotations‘, ‘instances_‘
                                      + self.set_name + ‘.json‘))
        self.image_ids = self.coco.getImgIds()
        self.load_classes()
    
    def load_classes(self):
        categories = self.coco.loadCats(self.coco.getCatIds())
        categories.sort(key=lambda x: x[‘id‘])
        
        # coco ids is not from 1, and not continue
        # make a new index from 0 to 79, continuely
        
        # classes:             {names:      new_index}
        # coco_labels:         {new_index:  coco_index}
        # coco_labels_inverse: {coco_index: new_index}
        self.classes, self.coco_labels, self.coco_labels_inverse = {}, {}, {}
        for c in categories:
            self.coco_labels[len(self.classes)] = c[‘id‘]
            self.coco_labels_inverse[c[‘id‘]]   = len(self.classes)
            self.classes[c[‘name‘]] = len(self.classes)
        
        # labels:              {new_index:  names}
        self.labels = {}
        for k, v in self.classes.items():
            self.labels[v] = k

    def __len__(self):
        return len(self.image_ids)            
    
    def __getitem__(self, index):
        img = self.load_image(index)
        ann = self.load_anns(index)
        sample = {‘img‘:img, ‘ann‘: ann}
        
        if self.transform:
            sample = self.transform(sample)
        return sample
    
    def load_image(self, index):
        image_info = self.coco.loadImgs(self.image_ids[index])[0]
        imgpath       =  os.path.join(self.rootdir, ‘images‘, self.set_name, 
                                   image_info[‘file_name‘])
        
        img = skimage.io.imread(imgpath)
        return img.astype(np.float32) / 255.0
    
    def load_anns(self, index):
        annotation_ids = self.coco.getAnnIds(self.image_ids[index], iscrowd=False)
        # anns is num_anns x 5, (x1, x2, y1, y2, new_idx)
        anns = np.zeros((0, 5))
        
        # skip the image without annoations
        if len(annotation_ids) == 0:
            return anns
        
        coco_anns = self.coco.loadAnns(annotation_ids)
        for a in coco_anns:
            # skip the annotations with width or height < 1
            if a[‘bbox‘][2] < 1 or a[‘bbox‘][3] < 1:
                continue
            
            ann = np.zeros((1, 5))
            ann[0, :4] = a[‘bbox‘]
            ann[0, 4]  = self.coco_labels_inverse[a[‘category_id‘]]
            anns = np.append(anns, ann, axis=0)
        
        # (x1, y1, width, height) --> (x1, y1, x2, y2)
        anns[:, 2] += anns[:, 0]
        anns[:, 3] += anns[:, 1]
        
        return anns
    
    def image_aspect_ratio(self, index):
        image = self.coco.loadImgs(self.image_ids[index])[0]
        return float(image[‘width‘]) / float(image[‘height‘])