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 = 0for 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 + 1file_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_valuelbl = 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] = namelbl_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()