图像分割之json转为mask图像

import json
import os
import os.path as osp
import base64
import warnings
from PIL import Image

# import PIL.Image
import yaml
from labelme import utils
import cv2
import numpy as np
from skimage import img_as_ubyte


def main(json_file):
    warnings.warn("This script is aimed to demonstrate how to convert the\n"
                  "JSON file to a single image dataset, and not to handle\n"
                  "multiple JSON files to generate a real-use dataset.")
    # json文件加入列表
    list_path = []
    for json_single in os.listdir(json_file):
        if json_single.endswith('.json'):
            list_path.append(json_single)

    for i in range(0, len(list_path)):
        # 挨个读取json文件
        path = os.path.join(json_file, list_path[i])
        if os.path.isfile(path):
            data = json.load(open(path,encoding='utf-8'))
            img = utils.img_b64_to_arr(data['imageData'])
            print('imagedata: ', data['imageData'])
            print('datashape: ', data['shapes'])
            lbl, lbl_names = utils.labelme_shapes_to_label(img.shape, data['shapes'])

            captions = ['%d: %s' % (l, name) for l, name in enumerate(lbl_names)]

            lbl_viz = utils.draw_label(lbl, img, captions)
            # out_dir = osp.basename(path).replace('.', '_')
            out_dir = osp.basename(path).split('.json')[0]
            save_file_name = out_dir
            # out_dir = osp.join(osp.dirname(path), out_dir)

            if not osp.exists(json_file + 'mask'):
                os.mkdir(json_file + 'mask')
            maskdir = json_file + 'mask'

            if not osp.exists(json_file + 'mask_viz'):
                os.mkdir(json_file + 'mask_viz')
            maskvizdir = json_file + 'mask_viz'
            out_dir1 = maskdir

            # 保存
            Image.fromarray(lbl).save(out_dir1 + '/' + save_file_name + '.png')
            Image.fromarray(lbl_viz).save(maskvizdir + '/' + save_file_name + '_label_viz.png')

            with open(osp.join(out_dir1, 'label_names.txt'), 'w') as f:
                for lbl_name in lbl_names:
                    f.write(lbl_name + '\n')
            warnings.warn('info.yaml is being replaced by label_names.txt')
            info = dict(label_names=lbl_names)
            with open(osp.join(out_dir1, 'info.yaml'), 'w') as f:
                yaml.safe_dump(info, f, default_flow_style=False)
            print('Saved to: %s' % out_dir1)


if __name__ == '__main__':
    json_file = r"F:\个人实验代码\study\test\json"
    main(json_file)