import gradio as gr
from PIL import Image
import numpy as np


def draw_mask(img_info):
    img = img_info['image']
    img_mask = img_info['mask']
    img = img.convert('RGB')
    img_array= np.array(img)

    img_mask_array = np.array(img_mask.convert('RGB'))

    img_mask = img_mask.convert('L')
    blue_img_mask = img_mask.convert('RGB')
    blue_img_mask_array = np.array(blue_img_mask)
    blue_img_mask_array[:,:,2:] = 0
    blue_img = Image.fromarray(blue_img_mask_array)

    print(img_mask_array.shape)
    # print(img_mask_array.max()) #255
    img_mask_array = img_mask_array//255

    blend_img = (1-img_mask_array)*img_array + (img_mask_array)*(blue_img_mask_array)

    blend_img = Image.fromarray(blend_img)

    unmask_img = Image.fromarray((1-img_mask_array)*img_array)
    
    w, h = img_mask.size
    return [img, blend_img]



def print_slider(n_input):

    return f'{n_input}'


with gr.Blocks() as demo:
    gr.Markdown("Flip text or image files using this demo.")
    with gr.Tabs():
        with gr.TabItem("Flip Text"):
            # text_input = gr.Textbox()
            # text_output = gr.Textbox()
            # with gr.Column():
            #     n_input_btn = gr.Button()
            img_mask = gr.Image(value=None, source='upload', interactive=True,
            type='pil', tool='sketch', image_mode='RGBA')
            text_button = gr.Button("test")
            
            image_input = gr.Gallery().style(grid=[2,2])
    # n_input_btn.click(print_slider, inputs=n_input, outputs=text_output)
    text_button.click(draw_mask, inputs=[img_mask], outputs=[image_input])
    # img_mask.change(print_shpe, inputs=[img_mask], outputs=text_output)

    
demo.launch()