如何使用COCO数据集,注意事项

COCO数据集可用来训练目标检测,分类,实例分割等。

下面简单说下如何使用这个数据集,
数据集下载可用如下的代码进行,以2017为例。

# Download the image data.
cd ./images
echo "Downloading MSCOCO train images ..."
curl -LO http://images.cocodataset.org/zips/train2017.zip
echo "Downloading MSCOCO val images ..."
curl -LO http://images.cocodataset.org/zips/val2017.zip

cd ../
if [ ! -d annotations ]
  then
    mkdir -p ./annotations
fi

# Download the annotation data.
cd ./annotations
echo "Downloading MSCOCO train/val annotations ..."
curl -LO http://images.cocodataset.org/annotations/annotations_trainval2017.zip
echo "Finished downloading. Now extracting ..."

# Unzip data
echo "Extracting train images ..."
unzip -qqjd ../images ../images/train2017.zip
echo "Extracting val images ..."
unzip -qqjd ../images ../images/val2017.zip
echo "Extracting annotations ..."
unzip -qqd .. ./annotations_trainval2017.zip

会得到这2个文件夹

在这里插入图片描述
annotation文件夹下有各种json文件,记下它们的路径,后面要用到。

下面说下数据集处理的主要流程,主要用COCO API来提取数据,这个API是conda自带的,不需要另外安装
用的时候import就行了

from pycocotools.coco import COCO

刚才说了annotation文件夹下有json文件,假设这个json文件的路径为ann_path(包括json)
把这个路径传给COCO API

self.coco_api = COCO(ann_path)

然后你就能用这个coco_api提取各种数据,
比如提取类别转为label

self.cat_ids = sorted(self.coco_api.getCatIds())
self.cat2label = {cat_id: i for i, cat_id in enumerate(self.cat_ids)}

提取类别的名称

self.cats = self.coco_api.loadCats(self.cat_ids)
self.class_names = [cat["name"] for cat in self.cats]

提取img_id, 每个id对应一个img文件名,还能根据这个img_id提取对应的annotation,
把img_id对应的信息全部放到img_info里面

self.img_ids = sorted(self.coco_api.imgs.keys()) #这个是全部图片的img_id,需要的时候根据idx提取一个
img_info = self.coco_api.loadImgs(self.img_ids)

所以后面能用img_info干什么呢。
从pytorch的__getitem__函数说起吧,
getitem函数会传入一个idx, 提取这个idx对应的图片和annotation.

我们可以根据上面得到的全部图片的img_id, 也就是self.img_ids, 得到idx对应的img_id
再根据这个img_id得到这个图片对应的annotation

img_id = self.img_ids[idx]
ann_ids = self.coco_api.getAnnIds([img_id])
anns = self.coco_api.loadAnns(ann_ids)

得到图片的文件名,于是可以读取图片

file_name = self.coco_api.loadImgs(img_id)[0]['file_name']
if file_name.startswith('COCO'):
    file_name = file_name.split('_')[-1]

path = osp.join(self.root, file_name)
assert osp.exists(path), 'Image path does not exist: {}'.format(path)

img = cv2.imread(path)

下面说如何获得目标框,类别和分割mask
一个图片可能对应多个目标框,所以一个img_id提取出的anns里面可能有多个annotation,
可以用一个循环把它们读出来

for ann in anns:
    if ann.get("ignore", False):
        continue
    x1, y1, w, h = ann["bbox"]  #目标框是(x,y,w,h)形式
    if ann["area"] <= 0 or w < 1 or h < 1:
        continue
    if ann["category_id"] not in self.cat_ids:
        continue
    bbox = [x1, y1, x1 + w, y1 + h]  #有需要的话转为(x1,y1,x2,y2)形式
    
    gt_bboxes.append(bbox)
    gt_labels.append(self.cat2label[ann["category_id"]])
    
    gt_masks.append(self.coco_api.annToMask(ann).reshape(-1))  #提取分割mask
    
    gt_keypoints.append(ann["keypoints"])  #提取特征点

需要注意的是一个图片里面可能没有目标框,也就是annotation为[ ],
这种情况就不要把它当作训练图片了,需要重新采样,
所以需要这样一个流程

    def __getitem__(self, idx):
        
        while True:
              data = self.get_train_data(idx) #annotation为空时要返回None
              if data is None:
                  idx = self.get_another_id()
                  continue
              return data

    def get_another_id(self):
        return np.random.random_integers(0, len(self.data_info) - 1)

在预处理中,如果对图片进行各种变换,不要忘了同时要对目标框,分割mask做同样的变换

示例:

原图
在这里插入图片描述

对图像做翻转,scale,crop变换后

在这里插入图片描述

那么对mask, 目标框也要做同样的变换。

在这里插入图片描述

还有一种情况,annotation不为空,数据是有的,
但是,
在各种预处理变换后,目标没了,这个情况也是要重新采样的

示例:
在这里插入图片描述
它的目标比较小,在左上角的地方,很容易就没了,
比如做如下的预处理变换,

在这里插入图片描述

那怎么判断目标是不是没了呢,
一般来说,目标框经过变换后,会出现x1=x2, 或者y1=y2的情况,
也就是说这个box的w或者h为0

在这里插入图片描述
只需要把w和h很小的box滤掉即可,
滤掉后如果为空,就返回None, 让getitem函数重新采样。

w = meta_data["gt_bboxes"][:, 2] - meta_data["gt_bboxes"][:, 0]
h = meta_data["gt_bboxes"][:, 3] - meta_data["gt_bboxes"][:, 1]

keep = (w > self.discard_box_width) * (h > self.discard_box_height)
meta_data["gt_masks"] = torch.FloatTensor(meta_data["gt_masks"][keep])
meta_data["gt_bboxes"] = meta_data["gt_bboxes"][keep]
meta_data["gt_labels"] = meta_data["gt_labels"][keep]

if meta["gt_bboxes"].size == 0:
    return None 
  • 0
    点赞
  • 9
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

蓝羽飞鸟

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值