对于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信息保存:
# cv2.imwrite(seg_output_path, mask)
save_colored_mask(mask, seg_output_path)
保存方式有两种,一种为灰度模式、一种为调色板模式,其区别与各自的用法可以参考我的另一篇博客:语义分割标签——mask的读取与保存_读取mask标签-CSDN博客。
完整代码如下:
"""
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)
最终,保存的结果如下(我这里用的是调色板模式,如果用灰度图模式,则保存的为黑白图像):