对于coco2017数据集，所有标注信息都保存在json文件里。有时候，我们需要从中提取语义分割的mask并保存为png图片，以和原图一一对应。因此，这里提供一个脚本进行转换。
一、coco2017数据集结构
总的结构如下：
├─annotations
└─images
├─train2017
└─val2017其中，images存放的是训练集、验证集的图片原图；annotations中存放的是标注文件：
2017/09/02 03:04 91,865,115 captions_train2017.json
2017/09/02 03:04 3,872,473 captions_val2017.json
2017/09/02 03:02 469,785,474 instances_train2017.json
2017/09/02 03:02 19,987,840 instances_val2017.json
2017/09/02 03:04 238,884,731 person_keypoints_train2017.json
2017/09/02 03:04 10,020,657 person_keypoints_val2017.json对于语义分割，我们需要的是instances_train2017.json、instances_val2017.json这两个文件。
二、从json文件中提取mask信息
我们需要使用pycocotools这个工具，其可以解析coco的json文件，并从中提取相应信息。
首先，需要指定一个json文件，如instances_train2017.json，并解析其中的类别ID、图像ID：
coco = COCO(annotation_file)
catIds = coco.getCatIds()
imgIds = coco.getImgIds()然后，从imgIds中依次取得一个图像元素，并获取该图像中的标注信息：
img = coco.loadImgs(imgId)[0]
annIds = coco.getAnnIds(imgIds=img['id'], catIds=catIds, iscrowd=None)
anns = coco.loadAnns(annIds)由于有些图像没有分割信息，所以需要判断annIds的长度，大于0则继续提取mask信息：
mask = coco.annToMask(anns[0]) * anns[0]['category_id']
for i in range(len(anns) - 1):
mask += coco.annToMask(anns[i + 1]) * anns[i + 1]['category_id']上述代码中，每个mask都乘了一个anns[i]['category_id']，这是为了使每个类别在标注文件中对应的值不同。有些博客和GitHub上的代码没有这个，所以得出来的mask没有区别度，并不能用于训练分割模型。
最后，将mask信息保存：
扫描二维码关注公众号，回复：
12566395 查看本文章
# cv2.imwrite(seg_output_path, mask)
save_colored_mask(mask, seg_output_path)保存方式有两种，一种为灰度模式、一种为调色板模式，其区别与各自的用法可以参考我的另一篇博客：https://blog.csdn.net/oYeZhou/article/details/111934432。
完整代码如下：
"""
get semantic segmentation annotations from coco data set.
"""
from PIL import Image
import imgviz
import argparse
import os
import tqdm
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()
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']
for i in range(len(anns) - 1):
mask += coco.annToMask(anns[i + 1]) * anns[i + 1]['category_id']
img_origin_path = os.path.join(args.input_dir, 'images', 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)
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)
最终，保存的结果如下（我这里用的是调色板模式，如果用灰度图模式，则保存的为黑白图像）：
