import gradio as gr
import cv2
from skimage import segmentation, color
import numpy as np

# 功能 1: 影像分割
def segment_image(image, n_segments, compactness):
    segments = segmentation.slic(image, n_segments=n_segments, compactness=compactness)
    segmented_image = color.label2rgb(segments, image, kind='avg')
    return (segmented_image * 255).astype(np.uint8)

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

# 功能 3: 遮罩生成
def generate_mask(image, threshold1, threshold2):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, threshold1, threshold2)
    mask = cv2.dilate(edges, None, iterations=2)
    mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)[1]
    return mask

# 功能 4: 影像修復
def inpaint_image(image, mask, method):
    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    inpaint_method = cv2.INPAINT_TELEA if method == "Telea" else cv2.INPAINT_NS
    inpainted = cv2.inpaint(image, mask, 3, inpaint_method)
    return inpainted


# Gradio 界面設計
with gr.Blocks() as app:
    gr.Markdown("# 🌟 影像處理工具")
    gr.Markdown("使用本工具進行多種影像處理，如影像分割、邊緣檢測、遮罩生成及影像修復。")

    # 圖片輸入與範例圖片
    with gr.Row():
        gr.Markdown("### 範例圖片")
        gr.Examples(
            examples=[
                ["example1.jpg"], 
                ["example2.jpg"], 
                ["example3.jpg"],
                ["image (2).webp"],
                ["image (3).webp"],
                ["image (4).webp"]
            ],
            inputs=[gr.Image(type="numpy", label="上傳圖片")],
            outputs=None,
            label="選擇範例圖片"
        )

    # 各功能分區
    with gr.Tab("影像分割"):
        gr.Markdown("### 影像分割")
        input_image = gr.Image(type="numpy", label="上傳圖片")
        n_segments = gr.Slider(10, 500, value=100, label="分割區域數量")
        compactness = gr.Slider(0.1, 10.0, value=1.0, label="緊湊性")
        segmented_output = gr.Image(label="分割結果")
        gr.Button("執行").click(segment_image, inputs=[input_image, n_segments, compactness], outputs=segmented_output)

    with gr.Tab("邊緣檢測"):
        gr.Markdown("### 邊緣檢測")
        input_image = gr.Image(type="numpy", label="上傳圖片")
        threshold1 = gr.Slider(0, 255, value=50, label="閾值1")
        threshold2 = gr.Slider(0, 255, value=150, label="閾值2")
        edge_output = gr.Image(label="邊緣檢測結果")
        gr.Button("執行").click(edge_detection, inputs=[input_image, threshold1, threshold2], outputs=edge_output)

    with gr.Tab("遮罩生成"):
        gr.Markdown("### 遮罩生成")
        input_image = gr.Image(type="numpy", label="上傳圖片")
        threshold1 = gr.Slider(0, 255, value=50, label="邊緣檢測閾值1")
        threshold2 = gr.Slider(0, 255, value=150, label="邊緣檢測閾值2")
        generated_mask = gr.Image(label="生成的遮罩")
        gr.Button("生成遮罩").click(generate_mask, inputs=[input_image, threshold1, threshold2], outputs=generated_mask)

    with gr.Tab("影像修復"):
        gr.Markdown("### 影像修復")
        input_image = gr.Image(type="numpy", label="上傳圖片")
        input_mask = gr.Image(type="numpy", label="選擇遮罩（生成或手動上傳）")
        method = gr.Radio(["Telea", "Navier-Stokes"], value="Telea", label="修復演算法")
        inpainted_output = gr.Image(label="修復結果")
        gr.Button("執行修復").click(inpaint_image, inputs=[input_image, input_mask, method], outputs=inpainted_output)


app.launch()