import base64
from io import BytesIO
import cv2
from PIL import Image
import numpy as np
import gradio as gr


def encode(image):
    if image is None:
        raise gr.Error('No input image')

    buffer = BytesIO() 
    image.save(buffer, format='PNG')
    img_str = base64.b64encode(buffer.getvalue()).decode("utf-8")
    return img_str

def decode(image_str):
    if image_str is None:
        raise gr.Error('No input image string')

    image_byte = base64.b64decode(image_str)
    image = cv2.imdecode(np.frombuffer(image_byte, np.uint8), cv2.IMREAD_UNCHANGED)
    image = Image.fromarray(image[:,:,::-1])
    return image
    
with gr.Blocks() as demo:
    with gr.Tab(label='Encoding'):
        with gr.Row():
            with gr.Column():
                image = gr.Image(type='pil', interactive=True)
                btn_enc = gr.Button()
            output_image_str = gr.Textbox(label='image string', interactive=False)
    with gr.Tab('Decoding'):
        with gr.Row():
            with gr.Column():
                image_str = gr.Textbox(label='image string', interactive=True)
                btn_dec = gr.Button()
            output_image = gr.Image(type='pil', interactive=False)
            

    btn_enc.click(
        fn=encode,
        inputs=[image],
        outputs=[output_image_str],
    )

    btn_dec.click(
        fn=decode,
        inputs=[image_str],
        outputs=[output_image],
    )

demo.launch(show_error=True)