import base64
import json
import os
import os.path as osp

import numpy as np
import PIL.Image
from labelme import utils


def json2mask(root, classes)
    before_path = os.path.join(root, 'json') # labelme files
    jpgs_path = os.path.join(root,   'jpgs') # images
    pngs_path = os.path.join(root, 'masks')  # annotations
    if not os.path.exists(jpgs_path):
        os.makedirs(jpgs_path)
    if not os.path.exists(pngs_path):
        os.makedirs(pngs_path)

    path = before_path
    for file in os.listdir(path):
        if file.endswith('json'):
            data = json.load(open(os.path.join(path, file)))

            if data['imageData']:
                imageData = data['imageData']
            else:
                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
                with open(imagePath, 'rb') as f:
                    imageData = f.read()
                    imageData = base64.b64encode(imageData).decode('utf-8')

            img = utils.img_b64_to_arr(imageData)
            label_name_to_value = {classes[0]: 0, classes[2]:0, classes[1]: 1 }

            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)

            PIL.Image.fromarray(img).save(osp.join(jpgs_path, file.split(".")[0] + '.jpg'))

            utils.lblsave(osp.join(pngs_path, file.split(".")[0] + '.png'), lbl)

if __name__ == '__main__':
    root = 'J:/dataset_panicle/2023/only_plant/images'
    classes = ["_background_",  "panicle", "other"]
    json2mask(root,classes)