# app.py

import gradio as gr
import torch
import numpy as np
import torch.nn.functional as F
from PIL import Image
from huggingface_hub import hf_hub_download
from inference import load_model, predict

# 1) Descarga automática de los pesos desde el Model Hub
model_path = hf_hub_download(
    repo_id="vncgabriel/instancia-segmentation-model",
    filename="pytorch_model.bin",
    repo_type="model",
)

# 2) Carga el modelo (usa GPU si está disponible)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = load_model(model_path, device)
model.eval()

def segmentar_imagen(image: Image.Image):
    """
    Recibe una PIL Image y devuelve:
      1) overlay: imagen original con máscara semitransparente en rojo
      2) pure: máscara pura en amarillo sobre fondo negro
    """
    # Preprocesado: RGB -> numpy [H,W,3] -> tensor [1,3,H,W]
    img = image.convert("RGB")
    arr = np.array(img, dtype=np.float32) / 255.0
    tensor = torch.from_numpy(arr).permute(2,0,1).unsqueeze(0).to(device)

    # Padding a múltiplos de 32 (5 downsamples)
    _,_,H,W = tensor.shape
    pad_h = (32 - H % 32) % 32
    pad_w = (32 - W % 32) % 32
    tensor_p = F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect")

    # Inferencia
    with torch.no_grad():
        mask_p = model(tensor_p)[0,0]   # [H+pad, W+pad]

    # Binariza y recorta al tamaño original
    mask_np = (mask_p.cpu().numpy() > 0.5).astype(np.uint8) * 255
    mask = mask_np[:H, :W]

    # --- Overlay en rojo ---
    overlay = img.convert("RGBA")
    mask_img = Image.fromarray(mask).convert("L")
    capa_roja = Image.new("RGBA", overlay.size, (255, 0, 0, 100))
    overlay.paste(capa_roja, mask=mask_img)

    # --- Máscara pura en amarillo ---
    pure = Image.new("RGB", (W, H), (0, 0, 0))
    capa_amarilla = Image.new("RGB", (W, H), (255, 255, 0))
    pure.paste(capa_amarilla, mask=mask_img)

    return overlay, pure

# 3) Interfaz Gradio
iface = gr.Interface(
    fn=segmentar_imagen,
    inputs=gr.Image(type="pil", label="Imagen de entrada"),
    outputs=[
        gr.Image(type="pil", label="Overlay en rojo"),
        gr.Image(type="pil", label="Máscara pura en amarillo"),
    ],
    title="Segmentación de Instancias (Overlay + Máscara)",
    description="Sube una imagen y obtén la segmentación de instancias: overlay rojo y máscara amarilla.",
    live=True,
)

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