labelme生成png

标签映射根据具体情况修改，生成的mask为彩色图。

import argparse
import glob
import os
import shutil

import labelme
import tqdm


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_dir', help='input directory', default="D:/Projects/UNet/Data/")
    parser.add_argument('--output_dir', help='output directory', default="D:/Projects/UNet/Dataset")
    args = parser.parse_args()

    if os.path.exists(args.output_dir):
        print('Output directory already exists: ', args.output_dir)
        shutil.rmtree(args.output_dir)
        # sys.exit(1)

    os.makedirs(args.output_dir)

    os.makedirs(os.path.join(args.output_dir, 'image'))
    os.makedirs(os.path.join(args.output_dir, 'mask'))
    class_name_to_id = {'none': 0, '0': 1, '2': 2, '3': 3, '4': 4}

    json_file_names = glob.glob(os.path.join(args.input_dir, '*.json'))
    image_names = []
    for file_name in tqdm.tqdm(json_file_names):
        # print('Generating test sample from:', file_name)
        label_file = labelme.LabelFile(filename=file_name)
        base = os.path.splitext(os.path.basename(file_name))[0]
        image_names.append(base + '.png')
        out_img_file = os.path.join(args.output_dir, 'image', base + '.png')
        out_mask_file = os.path.join(args.output_dir, 'mask', base + '.png')
        with open(out_img_file, 'wb') as f:
            f.write(label_file.imageData)
        img = labelme.utils.img_data_to_arr(label_file.imageData)
        # print(img.shape," ",class_name_to_id)
        lbl, _ = labelme.utils.shapes_to_label(
            img_shape=img.shape,
            shapes=label_file.shapes,
            label_name_to_value=class_name_to_id)
        labelme.utils.lblsave(out_mask_file, lbl)


if __name__ == '__main__':
    main()