pytorch coco 目标检测 DataLoader实现
pytorch实现目标检测目标检测算法首先要实现数据的读入,即实现Dataset
和DataLoader
两个类。
借助pycocotools
实现了CoCo2017用于目标检测数据的读取,并使用cv2
显示。
分析
使用cv2
显示读入数据,或者要送入到网络的数据应该有三个部分
- 图像,Nx3xHeight x Width
- BBs,NxMx4
- 类型,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和类型的具体信息。
要注意的事情有:
- python中图像的读入的通常是numpy的uint8数组,需要转换成float类型,并除以255以使最大值为1.0;
- coco数据中有80个类型,但是给的标签值最大为90,说明并不连续,需要设置新的标签,新的标签要从0到79,一定从0开始。
- 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‘])
原文地址:https://www.cnblogs.com/zi-wang/p/9972102.html