import gradio as gr
import cv2
import numpy as np

# 圖像分割功能
def image_segmentation(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    _, segmented = cv2.threshold(gray, 128, 255, cv2.THRESH_BINARY)
    return segmented

# 邊緣檢測功能
def edge_detection(image, threshold1, threshold2):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, threshold1, threshold2)
    return edges

# 圖像修復功能（Inpainting）
def image_inpainting(image, mask):
    if mask is None:
        return "請提供遮罩圖像以執行圖像修復。"
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY) if len(mask.shape) == 3 else mask
    inpainted = cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
    return inpainted

# 定義 Gradio 接口
def main_interface(image, task, param1=50, param2=150, mask=None):
    if task == "Image Segmentation":
        return image_segmentation(image)
    elif task == "Edge Detection":
        return edge_detection(image, param1, param2)
    elif task == "Image Inpainting":
        return image_inpainting(image, mask)
    else:
        return "Invalid Task or Missing Parameters!"

# 示例圖片
example_images = [
    ["example1.jpg", "測試圖片 1"],
    ["example2.jpg", "測試圖片 2"]
]

# UI 設計
with gr.Blocks() as app:
    gr.Markdown("# 電腦視覺功能展示應用程式")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="上傳圖片", type="numpy")
            task = gr.Radio(["Image Segmentation", "Edge Detection", "Image Inpainting"], label="選擇功能")
            param1 = gr.Slider(0, 200, value=50, step=1, label="參數 1 (Edge Detection Threshold 1)")
            param2 = gr.Slider(0, 200, value=150, step=1, label="參數 2 (Edge Detection Threshold 2)")
            mask_input = gr.Image(label="上傳遮罩 (僅適用於 Image Inpainting)", type="numpy", value=None)
            submit_button = gr.Button("執行")
        with gr.Column():
            output_image = gr.Image(label="輸出結果")
    
    gr.Examples(
        examples=example_images,
        inputs=input_image,
        label="示例圖片"
    )
    
    submit_button.click(main_interface, inputs=[input_image, task, param1, param2, mask_input], outputs=output_image)

# 啟動應用程式
if __name__ == "__main__":
    app.launch()