从COCO数据集标注获取实例分割掩码

· 本博客主要用来记录和分享经验，有问题欢迎大家友善讨论，转载请注明出处。

· 博主作为新人，水平有限，多谢理解。

前言

        COCO2017数据集中，用于分割任务的图片标注采用json文件存储，下面来提取出其中每幅图片的实例标注。

使用的文件：

代码

"""
get semantic or instance segmentation annotations from coco data set.
if semantic: use line35、38
if instance: use line36、39
"""
from PIL import Image
import imgviz
import argparse
import os
import tqdm
import shutil
import numpy as np
from pycocotools.coco import COCO

def save_colored_mask(mask, save_path):
    lbl_pil = Image.fromarray(mask.astype(np.uint8), mode="P")
    colormap = imgviz.label_colormap()
    lbl_pil.putpalette(colormap.flatten())
    lbl_pil.save(save_path)

def main(args):
    annotation_file = os.path.join(args.input_dir, 'annotations', 'instances_{}.json'.format(args.split))
    os.makedirs(os.path.join(args.input_dir, 'SegmentationClass'), exist_ok=True)
    os.makedirs(os.path.join(args.input_dir, 'JPEGImages'), exist_ok=True)
    coco = COCO(annotation_file)
    catIds = coco.getCatIds()
    imgIds = coco.getImgIds()
    empty_anno = []
    print("catIds len:{}, imgIds len:{}".format(len(catIds), len(imgIds)))
    for imgId in tqdm.tqdm(imgIds, ncols=100):
        img = coco.loadImgs(imgId)[0]
        annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
        anns = coco.loadAnns(annIds)
        if len(annIds) > 0:
            # mask = coco.annToMask(anns[0]) * anns[0]['category_id']
            mask = coco.annToMask(anns[0])
            for i in range(len(anns) - 1):
                # mask += coco.annToMask(anns[i + 1]) * anns[i + 1]['category_id']
                mask += coco.annToMask(anns[i + 1]) * (i + 2)
            img_origin_path = os.path.join(args.input_dir, args.split, img['file_name'])
            img_output_path = os.path.join(args.input_dir, 'JPEGImages', img['file_name'])
            seg_output_path = os.path.join(args.input_dir, 'SegmentationClass',
                                           img['file_name'].replace('.jpg', '.png'))
            shutil.copy(img_origin_path, img_output_path)
            save_colored_mask(mask, seg_output_path)
        else:
            empty_anno.append(imgId)
     print("No annotations images:", empty_anno)
     print("The number is ", len(empty_anno))

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_dir", default="../dataset/coco2017", type=str,
                        help="input dataset directory")
    parser.add_argument("--split", default="train2017", type=str,
                        help="train2017 or val2017")
    return parser.parse_args()

if __name__ == '__main__':
    args = get_args()
    main(args)

说明

1. COCO数据集中存在部分图片没有标注的情况，因为COCO指定的80个类不一定都在这些图片中出现。在代码中main函数里我添加了一个empty_anno列表，用于存储不含标注的图片名，最终输出该列表并给出共有多少张图片不含标注。

2. 参考1中该博主给出了语义分割技巧，在main函数的if语句内部计算掩码时，对每一个目标的掩码乘上category_id，每个相同类别目标的像素值相同，不同类别的目标像素值不同，这样前景目标就会按照类别分开。我这里针对实例分割，将category_id改为(i+2)，这样每个实例的像素值都不同。注意，不可以使用(i+1)，因为第一个目标我没有乘任何数，相当于乘了1，下面的for循环中i从0开始，若使用(i+1)则第二个目标也是乘1，导致前两个目标的像素值相同。

结果

train：包含1021张不含标注的图片；val：包含48张不含标注的图片。

参考

1.https://blog.csdn.net/oyezhou/article/details/111994155

感谢此博主的工作，给出了语义分割的代码，我的代码是基于他的代码修改的。

2.https://www.cnblogs.com/bob-jianfeng/p/11150821.html

同样感谢此博主的工作，给出了二值掩码的获取方法。