import gradio as gr
from skimage import io, color, feature
from skimage.segmentation import slic, mark_boundaries
from skimage.filters import threshold_otsu
import numpy as np
import cv2
import os

# 驗證影像格式
def validate_image(image):
    if not isinstance(image, np.ndarray):
        raise ValueError("Input is not a valid NumPy array.")
    if image.ndim == 2:  # 灰階影像
        image = np.stack([image] * 3, axis=-1)  # 擴展為 3 通道
    elif image.ndim == 3 and image.shape[2] != 3:
        raise ValueError("Image must have 3 channels (RGB).")
    return image

# 邊緣檢測
def edge_detection(image, low_threshold, high_threshold):
    gray = color.rgb2gray(image)
    edges = feature.canny(gray, sigma=low_threshold)
    return (edges * 255).astype(np.uint8)

# 影像分割
def image_segmentation(image, num_segments):
    segments = slic(image, n_segments=num_segments, start_label=1)
    segmented_image = mark_boundaries(image, segments)
    return (segmented_image * 255).astype(np.uint8)

# 修復影像
def image_inpainting(image, mask_size):
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    _, mask = cv2.threshold(gray, threshold_otsu(gray), 255, cv2.THRESH_BINARY_INV)
    mask = cv2.dilate(mask, None, iterations=mask_size)
    inpainted = cv2.inpaint(image, mask, inpaintRadius=3, flags=cv2.INPAINT_TELEA)
    return inpainted

# 處理影像
def process_image(task, image, param1, param2):
    try:
        image = validate_image(image)
        if task == "Edge Detection":
            return edge_detection(image, param1, param2)
        elif task == "Image Segmentation":
            return image_segmentation(image, int(param1))
        elif task == "Image Inpainting":
            return image_inpainting(image, int(param1))
    except Exception as e:
        print(f"Error: {e}")
        return np.zeros((100, 100, 3), dtype=np.uint8)

# Gradio 設定
tasks = ["Edge Detection", "Image Segmentation", "Image Inpainting"]
examples_dir = "examples"
examples = [
    ["Edge Detection", os.path.join(examples_dir, "desktop_Example3.png"), 2, 5],
    ["Image Segmentation", os.path.join(examples_dir, "desktop_Example3.png"), 100, 0],
    ["Image Inpainting", os.path.join(examples_dir, "desktop_Example3.png"), 5, 0],
]

interface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Dropdown(choices=tasks, label="Task"),
        gr.Image(type="numpy", label="Input Image"),
        gr.Slider(1, 10, value=1, label="Parameter 1 (e.g., Edge Sigma, Num Segments)"),
        gr.Slider(0, 10, value=1, label="Parameter 2 (e.g., High Threshold)"),
    ],
    outputs="image",
    examples=examples,
    title="Computer Vision Web App",
    description="Perform Edge Detection, Image Segmentation, or Inpainting on uploaded images.",
)

if __name__ == "__main__":
    interface.launch()