def test_shapes_to_label():
img, data = get_img_and_data()
label_name_to_value = {}
for shape in data['shapes']:
label_name = shape['label']
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
cls = shape_module.shapes_to_label(img.shape, data['shapes'],
label_name_to_value)
assert cls.shape == img.shape[:2]
def get_img_and_lbl():
img, data = get_img_and_data()
label_name_to_value = {"__background__": 0}
for shape in data["shapes"]:
label_name = shape["label"]
label_value = len(label_name_to_value)
label_name_to_value[label_name] = label_value
n_labels = max(label_name_to_value.values()) + 1
label_names = [None] * n_labels
for label_name, label_value in label_name_to_value.items():
label_names[label_value] = label_name
lbl, _ = shape_module.shapes_to_label(img.shape, data["shapes"],
label_name_to_value)
return img, lbl, label_names
def main():
parser = argparse.ArgumentParser(
formatter_class=argparse.ArgumentDefaultsHelpFormatter
)
parser.add_argument('input_dir', help='input annotated directory')
parser.add_argument('output_dir', help='output dataset directory')
parser.add_argument('--labels', help='labels file', required=True)
args = parser.parse_args()
if osp.exists(args.output_dir):
print('Output directory already exists:', args.output_dir)
sys.exit(1)
os.makedirs(args.output_dir)
os.makedirs(osp.join(args.output_dir, 'JPEGImages'))
os.makedirs(osp.join(args.output_dir, 'SegmentationClass'))
os.makedirs(osp.join(args.output_dir, 'SegmentationClassPNG'))
os.makedirs(osp.join(args.output_dir, 'SegmentationClassVisualization'))
print('Creating dataset:', args.output_dir)
class_names = []
class_name_to_id = {}
for i, line in enumerate(open(args.labels).readlines()):
class_id = i - 1 # starts with -1
class_name = line.strip()
class_name_to_id[class_name] = class_id
if class_id == -1:
assert class_name == '__ignore__'
continue
elif class_id == 0:
assert class_name == '_background_'
class_names.append(class_name)
class_names = tuple(class_names)
print('class_names:', class_names)
out_class_names_file = osp.join(args.output_dir, 'class_names.txt')
with open(out_class_names_file, 'w') as f:
f.writelines('\n'.join(class_names))
print('Saved class_names:', out_class_names_file)
#colormap1 = labelme.utils.label_colormap(255)
colormap1 = np.zeros((256, 3), dtype=float)
colormap1[0] = [0, 0, 0]
colormap1[1] = [255 ,0 ,0] #[127, 127, 127]
colormap1[2] = [191, 191, 191]
colormap1[3] = [255, 255, 255]
colormap1[4] = [64, 64, 64]
colormap1 = colormap1.astype(np.float32) / 255
print(colormap1)
for label_file in glob.glob(osp.join(args.input_dir, '*.json')):
print('Generating dataset from:', label_file)
with open(label_file) as f:
base = osp.splitext(osp.basename(label_file))[0]
out_img_file = osp.join(
args.output_dir, 'JPEGImages', base + '.jpg')
out_lbl_file = osp.join(
args.output_dir, 'SegmentationClass', base + '.npy')
out_png_file = osp.join(
args.output_dir, 'SegmentationClassPNG', base + '.png')
out_viz_file = osp.join(
args.output_dir,
'SegmentationClassVisualization',
base + '.jpg',
)
data = json.load(f)
img_file = osp.join(osp.dirname(label_file), data['imagePath'])
#img_name = str(label_file).split("/")[1].strip(".json")+".jpg"
#print(img_name)
#img_file = osp.join(osp.dirname(label_file), img_name)
#img_file = data['imagePath']
img = np.asarray(PIL.Image.open(img_file))
PIL.Image.fromarray(img).save(out_img_file)
lbl = shapes_to_label(
img_shape=img.shape,
shapes=data['shapes'],
label_name_to_value=class_name_to_id,
)
lblsave(out_png_file, lbl)
np.save(out_lbl_file, lbl)
viz = labelme.utils.draw_label(
lbl, img, class_names, colormap=colormap1)
PIL.Image.fromarray(viz).save(out_viz_file)
