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

examples = [
    '176788.jpg',
    '182862.jpg'
]
masks = [
    'mask1.jpg',
    'mask2.jpg'
]

def edge_detection(mode_selector, upload_input, image_selector, image_selected, threshold1, threshold2):
    if mode_selector == "Upload":
        image = upload_input
    else:
        image = image_selected
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, threshold1, threshold2)
    return edges

def mode_selector_changed(mode_selector, image_selector):
    if mode_selector == "Upload":
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=image_selector is not None, value=image_selector)

def mode_selector_changed_inpainting(mode_selector, image_selector, mask_selector):
    if mode_selector == "Upload":
        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
    else:
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=image_selector is not None, value=image_selector), gr.update(visible=True), gr.update(visible=image_selector is not None, value=mask_selector)

def close_selected_image(image_selected):
    if image_selected is None:
        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
    else:
        return gr.update(visible=True), gr.update(visible=True)
    
def close_selected_mask(mask_selected):
    if mask_selected is None:
        return gr.update(visible=True, value=None), gr.update(visible=False, value=None)
    else:
        return gr.update(visible=True), gr.update(visible=True)

def image_segmentation(mode_selector, upload_input, image_selector, image_selected):
    if mode_selector == "Upload":
        image = upload_input
    else:
        image = image_selected
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, segmented = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    return segmented

def image_inpainting(mode_selector, upload_input, upload_mask, image_selector, image_selected, mask_selector, mask_selected):  
    if mode_selector == "Upload":
        image = upload_input
        mask = upload_mask
    else:
        image = image_selected
        mask = mask_selected

    if mask.shape[:2] != image.shape[:2]:
        mask = cv2.resize(mask, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
    
    if len(mask.shape) == 3:
        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    mask = (mask > 0).astype(np.uint8) * 255
    inpainted = cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA)
    return inpainted

with gr.Blocks() as demo:
    with gr.Tab("Edge Detection"):
        mode_selector = gr.Radio(["Upload", "Example"], value="Upload", label="Upload an image or using example image")

        # Mode 1
        upload_input = gr.Image(type="numpy", label="Upload Image", visible=True)

        # Mode 2
        image_selector = gr.Radio(examples, label="Select an example image", visible=False)
        image_selected = gr.Image(None, label="Selected Image", visible=False, sources=[])

        threshold1 = gr.Slider(0, 255, value=50, step=1, label="Threshold 1")
        threshold2 = gr.Slider(0, 255, value=100, step=1, label="Threshold 2")

        output = gr.Image(type="numpy", label="Edge Image")
        inputs = [
            mode_selector,
            upload_input,
            image_selector,
            image_selected,
            threshold1,
            threshold2
        ]

        mode_selector.change(
            mode_selector_changed,
            inputs=[mode_selector, image_selector],
            outputs=[upload_input, image_selector, image_selected]
        )

        image_selector.change(
            mode_selector_changed,
            inputs=[mode_selector, image_selector],
            outputs=[upload_input, image_selector, image_selected]
        )

        image_selected.change(
            close_selected_image,
            inputs=[image_selected],
            outputs=[image_selector, image_selected]
        )

        gr.Interface(
            fn=edge_detection,
            inputs=inputs,
            outputs=output,
            description="Upload an image and adjust the thresholds to perform edge detection",
            flagging_mode='never'
        )
    with gr.Tab("Image Segmentation"):
        mode_selector = gr.Radio(["Upload", "Example"], value="Upload", label="Upload an image or using example image")

        # Mode 1
        upload_input = gr.Image(type="numpy", label="Upload Image", visible=True)
        # Mode 2
        image_selector = gr.Radio(examples, label="Select an example image", visible=False)
        image_selected = gr.Image(None, label="Selected Image", visible=False, sources=[])

        output = gr.Image(type="numpy", label="Edge Image")
        inputs = [
            mode_selector,
            upload_input,
            image_selector,
            image_selected,
        ]
        mode_selector.change(
            mode_selector_changed,
            inputs=[mode_selector, image_selector],
            outputs=[upload_input, image_selector, image_selected]
        )
        image_selector.change(
            mode_selector_changed,
            inputs=[mode_selector, image_selector],
            outputs=[upload_input, image_selector, image_selected]
        )
        image_selected.change(
            close_selected_image,
            inputs=[image_selected],
            outputs=[image_selector, image_selected]
        )
        gr.Interface(
            fn=image_segmentation,
            inputs=inputs,
            outputs=output,
            description="Upload an image and adjust the thresholds to perform edge detection",
            flagging_mode='never'
        )
    with gr.Tab("Image Inpainting"):
        mode_selector = gr.Radio(["Upload", "Example"], value="Upload", label="Upload an image or using example image")

        # Mode 1
        upload_input = gr.Image(type="numpy", label="Upload Image", visible=True)
        upload_mask = gr.Image(type="numpy", label="Upload Mask", visible=True)
        # Mode 2
        image_selector = gr.Radio(examples, label="Select an example image", visible=False)
        image_selected = gr.Image(None, label="Selected Image", visible=False, sources=[])
        mask_selector = gr.Radio(masks, label="Select an example mask", visible=False)
        mask_selected = gr.Image(None, label="Selected Mask", visible=False, sources=[])

        output = gr.Image(type="numpy", label="Edge Image")
        inputs = [
            mode_selector,
            upload_input,
            upload_mask,
            image_selector,
            image_selected,
            mask_selector,
            mask_selected
        ]
        mode_selector.change(
            mode_selector_changed_inpainting,
            inputs=[mode_selector, image_selector, mask_selector],
            outputs=[upload_input, upload_mask, image_selector, image_selected, mask_selector, mask_selected]
        )
        image_selector.change(
            mode_selector_changed,
            inputs=[mode_selector, image_selector],
            outputs=[upload_input, image_selector, image_selected]
        )
        mask_selector.change(
            mode_selector_changed,
            inputs=[mode_selector, mask_selector],
            outputs=[upload_mask, mask_selector, mask_selected]
        )
        image_selected.change(
            close_selected_image,
            inputs=[image_selected],
            outputs=[image_selector, image_selected]
        )
        mask_selected.change(
            close_selected_mask,
            inputs=[mask_selected],
            outputs=[mask_selector, mask_selected]
        )
        gr.Interface(
            fn=image_inpainting,
            inputs=inputs,
            outputs=output,
            description="Upload an image and adjust the thresholds to perform edge detection",
            flagging_mode='never'
        )

if __name__ == "__main__":
    demo.launch(allowed_paths=["/home/teng_aicv/Desktop/"])