import gradio as gr
from PIL import Image
import numpy as np
from transformers import CLIPSegProcessor, CLIPSegForImageSegmentation
from matplotlib.colors import to_rgb
import re
import cv2

# Load model
processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")

def parse_color(color_str):
    """
    Converts a color string (hex, name, or rgba(...)) to an RGB tuple.
    """
    try:
        if isinstance(color_str, str):
            if color_str.startswith("rgba("):
                # Extract the 3 RGB components
                numbers = list(map(float, re.findall(r"[\d.]+", color_str)))
                if len(numbers) >= 3:
                    r, g, b = numbers[:3]
                    return int(r), int(g), int(b)
            else:
                # Use named or hex color
                return tuple(int(255 * c) for c in to_rgb(color_str))
    except Exception:
        pass
    raise ValueError(f"Invalid color format: {color_str}. Use hex like '#ff0000', color name like 'red', or rgba format.")

def apply_mask(image: Image.Image, prompt: str, color: str) -> Image.Image:
    # Process the input image and prompt
    inputs = processor(text=prompt, images=image, return_tensors="pt")
    outputs = model(**inputs)
    preds = outputs.logits[0]

    # Get the binary mask from predictions
    mask = preds.sigmoid().detach().cpu().numpy()
    mask = (mask > 0.5).astype(np.uint8)

    # Convert image to RGBA
    image_np = np.array(image.convert("RGBA"))

    # Resize mask to match image size
    mask_resized = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]))
    mask_3d = np.stack([mask_resized] * 4, axis=-1)  # Extend mask to 3D

    # Convert the color string to an RGB tuple
    color_rgb = parse_color(color)
    overlay_color = np.array([*color_rgb, 128], dtype=np.uint8)  # RGBA with alpha 128

    # Create an overlay with the selected color
    overlay = np.zeros_like(image_np, dtype=np.uint8)
    overlay[:] = overlay_color

    # Apply the mask to the image
    masked_image = np.where(mask_3d == 1, overlay, image_np)
    return Image.fromarray(masked_image)

# Gradio Interface
iface = gr.Interface(
    fn=apply_mask,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Textbox(label="Segmentation Prompt", placeholder="e.g., helmet, road, sky"),
        gr.ColorPicker(label="Mask Color", value="#ff0000")
    ],
    outputs=gr.Image(type="pil", label="Segmented Image"),
    title="CLIPSeg Image Masking",
    description="Upload an image, input a prompt (e.g., 'person', 'sky'), and pick a mask color."
)

iface.launch()