JSON2Mask

import json
from PIL import Image
from PIL import ImageDraw
import PIL
import os
import numpy as np
from labelme import utils
import imgviz
import yaml


jsonfolder = './JSON/'
out_folder = 'labelme_json'

jsonfiles = os.listdir(jsonfolder)
for jsonfile in jsonfiles:
	name, ext = os.path.splitext(jsonfile)
	if ext != '.json':
		continue

	out_dir = os.path.join(out_folder, name)
	if not os.path.exists(out_dir):
		os.makedirs(out_dir)
	
	with open(jsonfolder + jsonfile,'r') as jf:
		labels = json.load(jf)
		shapes = labels['shapes']

		imageData = labels.get('imageData')

		# print(imageData.shape)
		img = utils.img_b64_to_arr(imageData)
		# print(img.shape)


		img_shape = img.shape
		# print('img_shape:', img_shape)

		ins_mask = np.zeros(img_shape[:2], dtype=np.uint8)
		ins_mask = PIL.Image.fromarray(ins_mask)
		label_names = ['__background__']
		
		objects = {}
		
		for j, region in enumerate(shapes):
			points = region['points']
			label_name = region['label']

			if label_name not in objects:
				objects[label_name] = []
				
			objects[label_name].append(points)

		j = 0
		for label_name, points_list in objects.items():
			for points in points_list:
				xy = [tuple(point) for point in points]
				draw = ImageDraw.Draw(ins_mask)
				draw.polygon(xy, outline=j+1, fill=j+1)
			label_names.append(label_name)
			j+=1

		lbl = np.array(ins_mask)

		save_img = PIL.Image.fromarray(img)

		# save_lbl = PIL.Image.fromarray(lbl)


		if img_shape[0] != 512 or img_shape[1] != 512:
			save_img = save_img.resize((512,512))
			# save_lbl = lbl.resize((512,512))
			lbl.reshpe(512,512)

		# PIL.Image.fromarray(img).save(os.path.join(out_dir, "img.png"))
		save_img.save(os.path.join(out_dir, "img.png"))

		utils.lblsave(os.path.join(out_dir, "label.png"), lbl)

		# with Image.open(filename) as image:

		# lbl_viz = imgviz.label2rgb(label=lbl, img=imgviz.asgray(img), label_names=label_names, loc="rb")
		# PIL.Image.fromarray(lbl_viz).save(os.path.join(out_dir, "label_viz.png"))

		with open(os.path.join(out_dir, "label_names.txt"), "w") as f:
			for lbl_name in label_names:
				f.write(lbl_name + "\n")

		info = dict(label_names = label_names)
		with open(out_dir + '/info.yaml','w') as f:
			yaml.safe_dump(info, f, default_flow_style=False)