import os
import io
import json
import numpy as np
from pycococreatortools import pycococreatortools
from PIL import Image
import base64
import cv2
from tqdm import tqdm
def img_tobyte(img_pil):
    
    ENCODING = 'utf-8'
    img_byte = io.BytesIO()
    img_pil.save(img_byte, format='PNG')
    binary_str2 = img_byte.getvalue()
    imageData = base64.b64encode(binary_str2)
    base64_string = imageData.decode(ENCODING)
    return base64_string

def mask2json(ROOT_DIR):
    Image_DIR = os.path.join(ROOT_DIR, "pngs")
    Label_DIR = os.path.join(ROOT_DIR, "pre")

    Label_files = os.listdir(Label_DIR)
    # classs_names
    class_names = ['_background_', 'panicle']
    for Label_filename in tqdm(Label_files):

        Json_output = {
            "version": "3.16.7",
            "flags": {},
            "fillColor": [255, 0, 0, 128],
            "lineColor": [0, 255, 0, 128],
            "imagePath": {},
            "shapes": [],
            "imageData": {}}
        name = Label_filename.split('.', 3)[0]
        name1 = name + '.png'
        Json_output["imagePath"] = name1

        image = Image.open(Image_DIR + '/' + name1)
        imageData = img_tobyte(image)
        Json_output["imageData"] = imageData
        binary_mask = np.asarray(np.array(Image.open(Label_DIR + '/' + Label_filename))).astype(np.uint8)
        mask_image = cv2.imread(Label_DIR + '/' + Label_filename, cv2.IMREAD_GRAYSCALE)
        temp_mask = np.asarray((mask_image != 0), dtype=np.uint8)

        segmentation = pycococreatortools.binary_mask_to_polygon(temp_mask, tolerance=2)
        for item in segmentation:
            if len(item) > 10:
                list1 = []
                for j in range(0, len(item), 2):
                    list1.append([item[j], item[j + 1]])
                # There is only one non-background class, so just use class_names[1]
                label = class_names[1]
                seg_info = {'points': list1, "fill_color": None, "line_color": None, "label": label,
                            "shape_type": "polygon", "flags": {}}
                Json_output["shapes"].append(seg_info)

        Json_output["imageHeight"] = binary_mask.shape[0]
        Json_output["imageWidth"] = binary_mask.shape[1]
        json_path = os.path.join(ROOT_DIR,'json')
        if not os.path.exists(json_path):
            os.makedirs(json_path)
        full_path = os.path.join(json_path, '{}.json'.format(name))
        with open(full_path, 'w') as output_json_file:
            json.dump(Json_output, output_json_file)


if __name__ == '__main__':

    ROOT_DIR = ''
    mask2json(ROOT_DIR)