labelme 标注的数据集转化为Mask-Rcnn适用的数据集

labelme 标注的数据集转化为Mask-Rcnn适用的数据集

食用步骤

1.labelme标注数据时，将生成的json文件和原图保存在一起

2.只需提供labelme生成的数据的文件夹，和maskrcnn的数据集文件夹，运行代码就会自动进行处理

3.代码会在提供的maskrcnn数据集文件夹下生成’cv2_mask’, ‘json’, ‘label’, ‘pic’,‘yaml’,'pic_and_mask’这几个文件夹

4.分别用于存储掩膜图片、json文件、标签txt文件、原图、yaml文件、带有掩膜的原图

5.根据自己需求，做对应的其他操作

import base64
import json
import os
import os.path as osp
import shutil
import PIL.Image
import yaml
from labelme.logger import logger
from labelme import utils
#将labelme生成的数据，转化为适用于maskrcnn的数据集。labelme标注数据时，将生成的json文件和原图保存在一起
#只需提供labelme生成的数据的文件夹，和maskrcnn的数据集文件夹，运行代码就会自动进行处理
#会在提供的maskrcnn数据集文件夹下生成'cv2_mask', 'json', 'label', 'pic','yaml','pic_and_mask'
#分别用于存储掩膜图片、json文件、标签txt文件、原图、yaml文件、带有掩膜的原图
def main():
    logger.warning('This script is aimed to demonstrate how to convert the'
                   'JSON file to a single image dataset, and not to handle'
                   'multiple JSON files to generate a real-use dataset.')

    labelme_json_file = 'your labelme data path'
    output_file='your maskrcnn dataset path'
    subdirs = ['cv2_mask', 'json', 'label', 'pic','yaml','pic_and_mask']
    int_file_name = 0
    for subdir in subdirs:
        # 组合得到完整的目录路径
        dir_path = os.path.join(output_file, subdir)

        os.makedirs(dir_path, exist_ok=True)


    for json_file in os.listdir(labelme_json_file):
        if json_file.endswith('.json'):
            int_file_name = int_file_name + 1
            file_name = str(int_file_name)
            data = json.load(open(labelme_json_file+'/'+json_file, encoding='utf-8'))

            imageData = data.get('imageData')
            #若不想用1，2，3来命名文件，可以使用下面代码，来获取文件本身的名字，用来命名
            file_name_with_extension = os.path.basename(labelme_json_file+'/'+json_file)#获取文件完整名字
            # file_names= os.path.splitext(file_name_with_extension)[0]#获取文件名字


            if not imageData:
                imagePath = os.path.join(os.path.dirname(json_file), data['imagePath'])
                with open(imagePath, 'rb',encoding='utf-8') 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}
            for shape in sorted(data['shapes'], key=lambda x: x['label']):
                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
            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)

            label_names = [None] * (max(label_name_to_value.values()) + 1)
            for name, value in label_name_to_value.items():
                label_names[value] = name
            lbl_viz = utils.draw_label(lbl, img, label_names)


            source_file_path = os.path.join(labelme_json_file, file_name_with_extension)
            target_file_path = os.path.join(output_file+'/json', file_name)
            shutil.copy2(source_file_path, target_file_path)
            PIL.Image.fromarray(img).save(osp.join(output_file+'/pic', file_name+'_img.bmp'))
            utils.lblsave(osp.join(output_file+'/cv2_mask', file_name+'_label.png'), lbl)
            PIL.Image.fromarray(lbl_viz).save(osp.join(output_file+'/pic_and_mask', file_name+'_label_viz.png'))

            with open(osp.join(output_file+'/label', file_name+'_label_names.txt'), 'w',encoding='utf-8') as f:
                for lbl_name in label_names:
                    f.write(lbl_name + '\n')

            logger.warning('info.yaml is being replaced by label_names.txt')
            info = dict(label_names=label_names)
            with open(osp.join(output_file+'/yaml', file_name+'_info.yaml'), 'w',encoding='utf-8') as f:
                yaml.safe_dump(info, f, default_flow_style=False)

            logger.info('Saved to: {}'.format(output_file))


if __name__ == '__main__':
    main()