使用labelme生成用于训练maskrcnn的工具代码

# -*- coding: utf-8 -*-
# ----------------------------
#!  Copyright(C) 2022
#   All right reserved.
#   文件名称：xxx.py
#   摘   要：xxx
#   当前版本:1.0
#   作   者：刘恩甫
#   完成日期：2022-x-x
# -----------------------------

import os
import cv2
import json
import numpy as np
from shutil import rmtree

if __name__ == '__main__':
    base_path=r'D:\liu_projects\label_data\img\\'
    save_path=r'./msk_data/'
    label_dict={'circle':1,'square':2}

    #创建文件夹
    if os.path.exists(save_path):
        rmtree(save_path)
    for k,v in label_dict.items():
        os.makedirs(save_path+str(v)+'/img/')
        os.makedirs(save_path+str(v)+'/label/')

    #解析数据，生成训练数据
    for path in os.listdir(base_path):
        name,ext=os.path.splitext(path)
        img=cv2.imread(base_path+name+'.jpg')
        json_data = json.load(open(base_path + name+'.json', 'r', encoding='utf-8'))
        w,h,c=img.shape

        label_img_list=[]
        for i in range(len(label_dict)):
            label_img_list.append(np.zeros((w,h),dtype=np.uint8))

        for j in range(len(json_data['shapes'])):
            label=json_data['shapes'][j]['label']
            point_array = np.array(json_data['shapes'][j]['points']).astype(np.int)
            cv2.fillConvexPoly(label_img_list[label_dict[label]-1], point_array, 255)

        for k,label_img in enumerate(label_img_list):
            cv2.imwrite(save_path+str(k+1)+'/img/'+name+'.jpg',img)
            cv2.imwrite(save_path+str(k+1)+'/label/'+name+'.png',label_img)

        # break