Python——coco格式图像分割数据集转mask

目前很多深度学习框架中的图像分割套件都使用image-mask格式的标签数据，所以为了方便使用写了该脚本进行转换。

单张coco转mask并显示

convert_coco2mask_single_show.py

from pycocotools.coco import COCO
import os
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
import cv2

def convert_coco2mask_show(image_id):
    img = coco.imgs[image_id]
    # loading annotations into memory...
    # Done (t=12.70s)
    # creating index...
    # index created!
    img
    # {'license': 2,
    #  'file_name': '000000000074.jpg',
    #  'coco_url': # 'http://images.cocodataset.org/train2017/000000000074.jpg',
    #  'height': 426,
    #  'width': 640,
    #  'date_captured': '2013-11-15 03:08:44',
    #  'flickr_url': # 'http://farm5.staticflickr.com/4087/5078192399_aaefdb5074_z.jpg# ',
    #  'id': 74}
    image = np.array(Image.open(os.path.join(img_dir, img['file_name'])))
    plt.imshow(image, interpolation='nearest')
    plt.show()
    plt.imshow(image)

    cat_ids = coco.getCatIds()
    anns_ids = coco.getAnnIds(imgIds=img['id'], catIds=cat_ids, iscrowd=None)
    anns = coco.loadAnns(anns_ids)
    coco.showAnns(anns)
    mask = coco.annToMask(anns[0])
    for i in range(len(anns)):
        mask += coco.annToMask(anns[i])
    plt.imshow(image, interpolation='nearest')
    plt.imshow(image)
    plt.imshow(mask)
    cv2.imwrite(os.path.join(save_dir, "mask{}.png".format(image_id)), mask)
    plt.show()

if __name__ == '__main__':
    Dataset_dir = "lane_coco_dataset"
    coco = COCO(os.path.join(Dataset_dir, 'Annotations/coco_info.json'))
    img_dir = os.path.join(Dataset_dir, 'Images')
    save_dir = os.path.join(Dataset_dir, "Mask")
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    image_id = 1
    convert_coco2mask_show(image_id)

批量coco转mask

• convert_coco2mask_batch.py

from pycocotools.coco import COCO
import os
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
import cv2

def convert_coco2mask():
    catIds = coco.getCatIds()
    imgIds = coco.getImgIds(catIds=catIds) # 获取所有图片的id
    print("Total images:", len(imgIds))
    for image_id in imgIds:
        img = coco.imgs[image_id]
        image = np.array(Image.open(os.path.join(img_dir, img['file_name'])))

        cat_ids = coco.getCatIds()
        anns_ids = coco.getAnnIds(imgIds=img['id'], catIds=cat_ids, iscrowd=None)
        anns = coco.loadAnns(anns_ids)
        coco.showAnns(anns)
        mask = coco.annToMask(anns[0])
        for i in range(len(anns)):
            mask += coco.annToMask(anns[i])
        cv2.imwrite(os.path.join(save_dir, "mask{}.png".format(image_id)), mask)

if __name__ == '__main__':
    Dataset_dir = "lane_coco_dataset"
    coco = COCO(os.path.join(Dataset_dir, 'Annotations/coco_info.json'))
    img_dir = os.path.join(Dataset_dir, 'Images')
    save_dir = os.path.join(Dataset_dir, "Mask")
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    convert_coco2mask()

导出后生成Mask文件夹：