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