labelme标注数据
使用labelme标注生成json后缀的文件，对于多个类别且每个类别有多个的数据，在由json文件生成掩膜图像时，要确保类别名称顺序正确。
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] # .jsonif 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_MAPlabelname = {}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)修改训练代码
主要是基于官方ballon训练数据修改load mask函数
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.sizenum_obj = np.max(img)mask = np.zeros([height, width, num_obj], dtype=np.uint8)#draw maskfor 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 = 1occlusion = np.logical_not(mask[:, :, -1]).astype(np.uint8)for i in range(count - 2, -1, -1):mask[:, :, i] = mask[:, :, i] * occlusionocclusion = np.logical_and(occlusion, np.logical_not(mask[:, :, i]))# 解析info.yamlyaml_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)