from __future__ import division, print_function, absolute_import
import numpy as np
import os
import json
from PIL import Image,ImageDraw
 
color_code={}
color_code['red']='#FF0000'
color_code['black']='#000000'
 
def get_each_mask(jsonpath):
    for file in os.listdir(jsonpath):
        json_file = jsonpath + file
        data = json.load(open(json_file))
        img = Image.new('RGB', [640, 640], color_code['black'])
        img1 = ImageDraw.Draw(img)
        label_idx_list = list(np.arange(len(data['shapes'])))
        for i in label_idx_list:
            if data['shapes'][i]['label'] == '1':
                xy = []
                for xy_tuple in data['shapes'][i]['points']:
                    xy += xy_tuple
                img1.polygon(xy, fill=color_code['red'], outline=color_code['red'])
                label_idx_list.remove(i)    
        img.save(r'/media/....../mask_thyroid/'+ os.path.splitext(file)[0] + '.png')
 
    for file in os.listdir(jsonpath):
        json_file = jsonpath + file
        data = json.load(open(json_file))
        img = Image.new('RGB', [640, 640], color_code['black'])
        img2 = ImageDraw.Draw(img)
        label_idx_list = list(np.arange(len(data['shapes'])))
        for i in label_idx_list:
            if data['shapes'][i]['label'] == '2':
                xy = []
                for xy_tuple in data['shapes'][i]['points']:
                    xy += xy_tuple
                img2.polygon(xy, fill=color_code['red'], outline=color_code['red'])
                label_idx_list.remove(i) 
        img.save(r'/media/....../mask_nodule/'+ os.path.splitext(file)[0] + '.png')
        
        
        
if __name__=='__main__':
    jsonpath=r'/media/....../json/'
    get_each_mask(jsonpath)