tensorflow+Mask RCNN训练labelme标注数据

    import argparse
    import json
    import os
    import os.path as osp
    import warnings
    import copy
    import numpy as np
    import PIL.Image
    from skimage import io
    import yaml
    from labelme import utils
    import codecs
    json_path = "/home/ubuntu/Documents/labelme/"
    out_dir1 = "/home/ubuntu/Documents/labelme_output/"
    list = os.listdir(json_path)
    for i in range(0, len(list)):
    json_file = os.path.join(json_path, list[i])
    filename = list[i][:-5]  # .json
    if os.path.isfile(json_file):
    #         with open(json_file, 'r',encoding="utf8", errors='ignore') as f:
    #            data = json.load(f)
        data = json.load(open(json_file))
        img = utils.image.img_b64_to_arr(data['imageData'])
        lbl, lbl_names = utils.shape.labelme_shapes_to_label(img.shape, data['shapes'])  # labelme_shapes_to_label
        # modify labels according to NAME_LABEL_MAP
        # modify labels according to NAME_LABEL_MAP
        labelname = {}
        captions = [] 
        for l, name in enumerate(lbl_names):
            for key_name in lbl_names:
                if lbl_names[key_name] == l:
                    captions.append(key_name)
                    labelname.update({key_name: l})
        lbl_viz = utils.draw.draw_label(lbl, img, captions)
        out_dir2 = osp.join(out_dir1,osp.dirname(list[i]))
        out_dir = osp.basename(list[i]).replace('.', '_')
        out_dir = osp.join(out_dir2,out_dir)
        if not osp.exists(out_dir):
            os.mkdir(out_dir)
    #       PIL.Image.fromarray(img).save(osp.join(out_dir, '{}.png'.format(filename)))
        PIL.Image.fromarray(lbl).save(osp.join(out_dir, '{}_gt.png'.format(filename)))
        PIL.Image.fromarray(lbl_viz).save(osp.join(out_dir, '{}_viz.png'.format(filename)))
        with open(osp.join(out_dir, 'label_names.txt'), 'w') as f:
            '''
            for l, name in enumerate(lbl_names):
                for key_name in lbl_names:
                    if lbl_names[key_name] == l:
                        f.write(key_name + '\n')
            '''
            for j in range(0, len(captions)):
                f.write(captions[j] + '\n')
        warnings.warn('info.yaml is being replaced by label_names.txt')
        #对字典按值从小到大排序
        info = dict(label_names=lbl_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)
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    
    AI写代码

    def load_mask(self, image_id):
        info = self.image_info[image_id]
        mask_path = info['path']        
        img = Image.open(mask_path)  # Image读取16位图像，CV2和Skimage.io读取８位图像。
        width, height = img.size
        num_obj = np.max(img)
        mask = np.zeros([height, width, num_obj], dtype=np.uint8)
        #draw mask
        for index in range(num_obj):
            for i in range(width):
                for j in range(height):
                    at_pixel = img.getpixel((i, j))
                    if at_pixel == index + 1:
                        mask[j, i, index] = 1
        # handle occlusion
        '# count is the number of object'
        count = 1
        occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)
        for i in range(count - 2, -1, -1):
            mask[:, :, i] = mask[:, :, i] * occlusion
            occlusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))
        # 解析info.yaml
        yaml_path = info['yaml_path']
        with open(yaml_path) as f:
            temp = yaml.load(f.read())
            labels = dict(temp['label_names'])
        labels_form = []
        class_ids = []
        name_dict = {"forklift": 1, "driver": 2}  ##0 is BG
        '''# 根据实际修改'''
        for l, name in enumerate(labels):
            for key_name in labels:
                if labels[key_name] == l:
                    if key_name.find("forklift") != -1:
                        labels_form.append("forklift")
                    elif key_name.find("driver") != -1:
                        labels_form.append("driver")
        # class name &class id
        [class_ids.append(name_dict[i]) for i in labels_form]
        class_ids = np.array(class_ids)
        return mask, class_ids.astype(np.int32)
    
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
      
    
    AI写代码