app.py

import base64
import hashlib
import io
import json
import traceback
import gradio as gr
import numpy as np
import torch
import cv2
from PIL import Image
from transformers import pipeline as hf_pipeline

# CPU-only: use more threads for better throughput
torch.set_num_threads(4)

sam_vit_pipeline = None

# ── Parametros sincronizados entre UI y backend ───────────────────────────────
PARAMS = {
    "pred_iou_thresh":        0.95,
    "stability_score_thresh": 0.5,
    "points_per_batch":       16,
    "min_mask_region_area":   4500,
    "box_nms_thresh":         0.8,
}


# ── Renderizado ───────────────────────────────────────────────────────────────
def _render_masks(imagen_rgb: Image.Image, masks: list) -> Image.Image:
    img_arr = np.array(imagen_rgb).copy()
    overlay = img_arr.copy()
    for i, mask in enumerate(masks):
        h = hashlib.md5(str(i).encode()).hexdigest()[:6]
        color = (int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16))
        overlay[np.array(mask) > 0] = color
    blended = cv2.addWeighted(img_arr, 0.5, overlay, 0.5, 0)
    return Image.fromarray(blended)


def _load_pipeline():
    global sam_vit_pipeline
    if sam_vit_pipeline is None:
        print("Cargando SAM ViT-Huge (CPU)...")
        sam_vit_pipeline = hf_pipeline(
            "mask-generation",
            model="facebook/sam-vit-huge",
            device=-1,
        )


# ── Segmentacion UI ───────────────────────────────────────────────────────────
def segmentar(
    imagen: Image.Image,
    pred_iou_thresh: float,
    stability_score_thresh: float,
    points_per_batch: int,
    min_mask_region_area: int,
    box_nms_thresh: float,
):
    global PARAMS
    if imagen is None:
        return None, "Sube una imagen para comenzar."

    PARAMS.update({
        "pred_iou_thresh":        float(pred_iou_thresh),
        "stability_score_thresh": float(stability_score_thresh),
        "points_per_batch":       int(points_per_batch),
        "min_mask_region_area":   int(min_mask_region_area),
        "box_nms_thresh":         float(box_nms_thresh),
    })

    _load_pipeline()
    imagen_rgb = imagen.convert("RGB")
    resultado = sam_vit_pipeline(
        imagen_rgb,
        points_per_batch=PARAMS["points_per_batch"],
        pred_iou_thresh=PARAMS["pred_iou_thresh"],
        stability_score_thresh=PARAMS["stability_score_thresh"],
        min_mask_region_area=PARAMS["min_mask_region_area"],
        box_nms_thresh=PARAMS["box_nms_thresh"],
    )
    if isinstance(resultado, list):
        resultado = resultado[0]

    masks = resultado.get("masks", [])
    if not masks:
        return imagen_rgb, "No se detectaron zonas."

    info = (
        f"UI: {len(masks)} zonas  |  "
        f"iou={PARAMS['pred_iou_thresh']}  stab={PARAMS['stability_score_thresh']}  "
        f"min_area={PARAMS['min_mask_region_area']}  "
        f"nms={PARAMS['box_nms_thresh']}  batch={PARAMS['points_per_batch']}"
    )
    return _render_masks(imagen_rgb, masks), info


# ── Endpoint para el backend Docker ──────────────────────────────────────────
def segment_for_backend(image_np: np.ndarray):
    """
    Llamado por el backend via gradio_client (api_name='/segment').
    Usa los mismos PARAMS que la UI — sincronizados al ultimo "Segmentar".
    Entrada : numpy uint8 H x W x 3.
    Salida  : (overlay_np, combined_json_str)
    """
    try:
        if image_np is None:
            empty = np.zeros((100, 100, 3), dtype=np.uint8)
            return empty, json.dumps({"masks": [], "label_map_b64": ""})

        _load_pipeline()
        pil_image = Image.fromarray(image_np.astype(np.uint8)).convert("RGB")
        h, w = image_np.shape[:2]

        resultado = sam_vit_pipeline(
            pil_image,
            points_per_batch=PARAMS["points_per_batch"],
            pred_iou_thresh=PARAMS["pred_iou_thresh"],
            stability_score_thresh=PARAMS["stability_score_thresh"],
            min_mask_region_area=PARAMS["min_mask_region_area"],
            box_nms_thresh=PARAMS["box_nms_thresh"],
        )
        if isinstance(resultado, list):
            resultado = resultado[0]

        all_masks_raw = resultado.get("masks", [])
        masks_bool = [np.array(m).astype(bool) for m in all_masks_raw]

        # Ordenar de mayor a menor area: grandes primero, pequenas al final para
        # que ventanas y detalles sobreescriban al muro en el label_map.
        masks_bool = sorted(masks_bool, key=lambda m: m.sum(), reverse=True)

        label_map = np.zeros((h, w), dtype=np.uint8)
        masks_out = []
        for i, mask in enumerate(masks_bool[:254], start=1):
            label_map[mask] = i
            area_ratio = float(mask.sum()) / max(1, h * w)
            ys, xs = np.where(mask)
            bbox = (
                [int(xs.min()), int(ys.min()), int(xs.max() - xs.min()), int(ys.max() - ys.min())]
                if len(ys) else [0, 0, 0, 0]
            )
            masks_out.append({
                "index": i,
                "surface": f"Zona {i}",
                "area_ratio": round(area_ratio, 4),
                "bbox_xywh": bbox,
            })

        pil_label = Image.fromarray(label_map, mode="L")
        buf = io.BytesIO()
        pil_label.save(buf, format="PNG")
        label_map_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")

        overlay_pil = _render_masks(pil_image, masks_bool)
        overlay_np = np.array(overlay_pil.convert("RGB"))

        combined = {
            "masks": masks_out,
            "label_map_b64": label_map_b64,
            "entorno": "cpu",
            "motor": "SAM Auto (CPU)",
            "params_used": dict(PARAMS),
        }
        return overlay_np, json.dumps(combined, ensure_ascii=False)

    except Exception:
        err = traceback.format_exc()
        empty = np.zeros((100, 100, 3), dtype=np.uint8)
        return empty, json.dumps({"error": err, "masks": [], "label_map_b64": ""})


# ── UI ────────────────────────────────────────────────────────────────────────
def crear_app():
    with gr.Blocks(title="SAM Auto - CPU") as demo:
        gr.Markdown("# Segmentacion Automatica - SAM ViT-Huge (CPU)")
        gr.Markdown(
            "SAM detecta todos los elementos de la imagen de forma automatica, "
            "sin necesidad de seleccionar zonas ni escribir prompts."
        )

        with gr.Row():
            imagen_entrada = gr.Image(type="pil", label="Foto del Espacio")
            imagen_salida  = gr.Image(label="Resultado")

        estado = gr.Markdown()
        boton  = gr.Button("Segmentar", variant="primary")

        with gr.Accordion("Parametros de segmentacion (sincronizados con el backend)", open=True):
            gr.Markdown(
                "> Los parametros que configures aqui se aplican tanto a la UI como al backend Docker. "
                "Haz clic en **Segmentar** para que el backend adopte los nuevos valores."
            )
            with gr.Row():
                sl_pred_iou = gr.Slider(
                    minimum=0.0, maximum=1.0, step=0.01, value=PARAMS["pred_iou_thresh"],
                    label="pred_iou_thresh  (↑ menos mascaras, mas limpias | HF default: 0.88)"
                )
                sl_stability = gr.Slider(
                    minimum=0.0, maximum=1.0, step=0.01, value=PARAMS["stability_score_thresh"],
                    label="stability_score_thresh  (↑ descarta zonas inestables | HF default: 0.95)"
                )
            with gr.Row():
                sl_batch = gr.Slider(
                    minimum=8, maximum=64, step=8, value=PARAMS["points_per_batch"],
                    label="points_per_batch  (en CPU mantener bajo, max recomendado: 16)"
                )
                sl_min_area = gr.Slider(
                    minimum=0, maximum=5000, step=100, value=PARAMS["min_mask_region_area"],
                    label="min_mask_region_area px  (↑ filtra zonas pequenas)"
                )
            with gr.Row():
                sl_nms = gr.Slider(
                    minimum=0.0, maximum=1.0, step=0.05, value=PARAMS["box_nms_thresh"],
                    label="box_nms_thresh  (↑ permite mas solapamiento entre mascaras)"
                )

        all_inputs = [imagen_entrada, sl_pred_iou, sl_stability, sl_batch, sl_min_area, sl_nms]
        boton.click(fn=segmentar, inputs=all_inputs, outputs=[imagen_salida, estado])
        imagen_entrada.upload(fn=segmentar, inputs=all_inputs, outputs=[imagen_salida, estado])

        # Endpoint oculto para el backend Docker
        _api_in   = gr.Image(type="numpy", label="backend_input",  visible=False)
        _api_over = gr.Image(type="numpy", label="backend_overlay", visible=False)
        _api_json = gr.Textbox(label="backend_json",                visible=False)
        _api_btn  = gr.Button(visible=False)
        _api_btn.click(
            fn=segment_for_backend,
            inputs=[_api_in],
            outputs=[_api_over, _api_json],
            api_name="segment",
        )

    return demo


demo = crear_app()
if __name__ == "__main__":
    demo.launch()