labelme转换特定标签mask

将labelme的特定标签转换成特定的mask

import json
import os
import os.path as osp
import cv2
from labelme import utils
import base64


def transJson(json_dir, spc_label=[], not_include=[], out_dir=None):
    json_file = json_dir
    if out_dir is None:
        out_dir = osp.basename(json_file).replace('.', '_')
        out_dir = osp.join(osp.dirname(json_file), out_dir)
    else:
        out_dir = out_dir
    if not osp.exists(out_dir):
        os.mkdir(out_dir)

    count = os.listdir(json_file)
    for i in range(0, len(count)):
        path = os.path.join(json_file, count[i])
        if os.path.isfile(path):
            data = json.load(open(path))

            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')
            img = utils.img_b64_to_arr(imageData)
            label_name_to_value = {'_background_': 0}

            # 这里执行过滤逻辑
            spec_list = []
            for item in data['shapes']:
                if item['label'] in spc_label:
                    spec_list.append(item)
                if item['label'] in not_include:
                    item['label'] = '_background_'
                    spec_list.append(item)
            data['shapes'] = spec_list

            # 这一段是为了获取从标签名到值的映射字典
            for shape in data['shapes']:
                label_name = shape['label']
                if label_name in label_name_to_value:
                    label_value = label_name_to_value[label_name]
                else:
                    label_value = len(label_name_to_value)
                    label_name_to_value[label_name] = label_value

            # label_values must be dense
            # 这一段得到标签值和标签名的数组
            label_values, label_names = [], []
            for ln, lv in sorted(label_name_to_value.items(), key=lambda x: x[1]):
                label_values.append(lv)
                label_names.append(ln)
            assert label_values == list(range(len(label_values)))

            # 这个函数获得一个shape，这个是生成mask的图像，传入了图像数据，data的shape， 还有label_name_to_value
            # 因此，只要控制label_name_to_value的映射，将特定的标签过滤或者指定特定的标签，就能实现生成指定的mask
            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)

            # captions = ['{}: {}'.format(lv, ln)
            #             for ln, lv in label_name_to_value.items()]
            # lbl_viz = utils.draw_label(lbl, img, captions)

            # PIL.Image.fromarray(img).save(osp.join(out_dir, 'img.png'))
            # PIL.Image.fromarray(lbl).save(osp.join(out_dir, 'label.png'))

            # 这个用来将上一个注释生成的包含mask的lbl信息保存成文件

					 # 如果报错的话，使用这句   utils.lblsave(osp.join(out_dir, count[i].replace('.json', '.png')), lbl)
					 # 代替下面这一局，这和labelme的版本有关。
            utils.lblsave(osp.join(out_dir, count[i].replace('.json', '.png')), lbl[0])

            label_png = cv2.imread(osp.join(out_dir, count[i].replace('.json', '.png')), 0)
            label_png[label_png > 0] = 255
            cv2.imwrite(osp.join(out_dir, count[i].replace('.json', '.png')), label_png)

            # PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, 'label_viz.png'))

            # with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
            #     for lbl_name in label_names:
            #         f.write(lbl_name + '\n')

            # warnings.warn('info.yaml is being replaced by label_names.txt')
            # info = dict(label_names=label_names)
            # with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
            #     yaml.safe_dump(info, f, default_flow_style=False)

            print('Saved to: %s' % out_dir)


if __name__ == '__main__':
    # 该文件夹只允许有json文件
    json_dir = 'E:\\CODE\\untitled1\\data\\LiverMask'
    # 包含的标签列表
    spec_labels = ['l']
    # 不包含的标签列表，用来去除和包含标签重合区域用
    not_include = ['t']

    # 这个地方用来存放生成的mask
    # out_dir = 'E:\\CODE\\untitled1\\data\\TrainFolder\\TumerMask2'
    out_dir = 'E:\\CODE\\untitled1\\data\\LiverMask'

    transJson(json_dir=json_dir, spc_label=spec_labels, not_include=not_include, out_dir=out_dir)
    print('finished!')