调用mmdetection框架代码进行COCO数据集加载

mmdetection目标检测框架提供了非常方便的数据集加载接口，只需要简单地修改几个参数，就可以配置不同的数据预处理方式。如果我们需要用COCO数据集进行除目标检测之外其他的计算机视觉方法研究，需要使用到COCO数据集，如何进行数据集加载呢？为了避免重复造轮子，本文提供了一种利用mmdetection框架代码的方法，只需要简单几行代码就可以实现数据加载。

from mmdet.datasets.coco import CocoDataset
from functools import partial
from mmcv.parallel import collate
from torch.utils.data import DataLoader

NumberCOCOImages = 500 # 决定batch_size有多大
data_root = '/home/dataset_coco/'# COCO数据集的路径
imgs_per_gpu = NumberCOCOImages
num_gpus = 1
workers_per_gpu = 1
batch_size = num_gpus * imgs_per_gpu
num_workers = num_gpus * workers_per_gpu

load_pipeline = [
    dict(type='LoadImageFromFile'),
    dict(type='Resize', img_scale=(224, 224), keep_ratio=False),#对原始图像进行缩放，全部变为224*224的分辨率
    dict(type='ImageToTensor', keys=['img']),#这个不能少，否则无法加载数据集
    dict(type='Collect', keys=['img'], meta_keys=('filename', 'ori_shape', 'img_shape', 'pad_shape', 'scale_factor'))] #这个不能少，否则无法加载数据集
    # 除了上述这些pipeline，还可以根绝需要参考官方配置文件进行增减
    
CocoDataset_rgb = CocoDataset(data_root + 'annotations/instances_val2017.json', load_pipeline,
                              data_root=data_root,
                              img_prefix=data_root + 'val2017/',
                              test_mode=True
                              )
                              
data_loader_rgb = DataLoader(
    CocoDataset_rgb,
    batch_size=batch_size,
    collate_fn=partial(collate, samples_per_gpu=imgs_per_gpu),
    num_workers=num_workers)
    
for i_rgb, data_rgb in enumerate(data_loader_rgb):
    img = data_rgb['img']
    print(img.shape)
    break

运行后的效果如下，可以看到，能够成功地得到最终的我们想要的tensor