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
import sys
from pathlib import Path

# ── ZeroGPU shim ─────────────────────────────────────────────────────────────
try:
    import spaces
except ImportError:
    class _DummySpaces:
        def GPU(self, fn):
            return fn
    spaces = _DummySpaces()

DEVICE = 0 if torch.cuda.is_available() else -1
sam_vit_pipeline = None

# ── Parametros sincronizados entre UI y backend ───────────────────────────────
# Estos valores se actualizan cada vez que el usuario corre "Segmentar" en la UI.
# segment_for_backend los lee para usar exactamente los mismos.
PARAMS = {
    "pred_iou_thresh":        0.95,
    "stability_score_thresh": 0.5,
    "points_per_batch":       32,
    "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...")
        sam_vit_pipeline = hf_pipeline(
            "mask-generation",
            model="facebook/sam-vit-huge",
            device=DEVICE,
        )


# ── Segmentacion UI ───────────────────────────────────────────────────────────
@spaces.GPU
@torch.no_grad()
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."

    # Sincronizar PARAMS con los sliders actuales
    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 ──────────────────────────────────────────
@spaces.GPU
@torch.no_grad()
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: las mascaras grandes se escriben primero
        # y las pequenas (ventanas, detalles) las sobreescriben — evita que el muro
        # tape a la ventana en el label_map.
        masks_bool = sorted(masks_bool, key=lambda m: m.sum(), reverse=True)

        # Label map: cada pixel contiene el indice de la mascara (1-based, max 254)
        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": "gpu",
            "motor": "SAM Auto (GPU - ZeroGPU)",
            "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": ""})


@spaces.GPU
@torch.no_grad()
def limpiar_mascara(mask: np.ndarray, area_minima: int = 1500) -> np.ndarray:
    """Elimina salpicaduras pequenas usando morfologia y connected components."""
    try:
        mask_uint8 = (mask.astype(np.uint8)) * 255
        kernel = np.ones((7, 7), np.uint8)
        mask_limpia = cv2.morphologyEx(mask_uint8, cv2.MORPH_OPEN, kernel)
        mask_limpia = cv2.morphologyEx(mask_limpia, cv2.MORPH_CLOSE, kernel)

        num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(mask_limpia, connectivity=8)
        mask_final = np.zeros_like(mask_limpia)

        if num_labels > 1:
            areas = stats[1:, cv2.CC_STAT_AREA]
            max_area = int(areas.max()) if areas.size else 0
            for i in range(1, num_labels):
                area = int(stats[i, cv2.CC_STAT_AREA])
                if area >= area_minima and (max_area == 0 or area >= (max_area * 0.05)):
                    mask_final[labels == i] = 255

        return (mask_final.astype(bool))
    except Exception:
        return mask


@spaces.GPU
@torch.no_grad()
def render_for_backend(
    image_np: np.ndarray,
    label_map_b64: str,
    mask_index: int = 1,
    texture_name: str | None = None,
    texture_b64: str | None = None,
    params_json: str = "{}",
):
    """
    Aplica la textura sobre la mascara especificada usando el pipeline del Generador de Texturas.
    Retorna (imagen_renderizada_np, json_str)
    """
    try:
        if image_np is None:
            empty = np.zeros((100, 100, 3), dtype=np.uint8)
            return empty, json.dumps({"error": "no_image"})

        # Preparar imagen
        pil_image = Image.fromarray(image_np.astype(np.uint8)).convert("RGB")
        image_rgb = np.array(pil_image)

        # Decodificar label_map
        if not label_map_b64:
            return image_rgb, json.dumps({"error": "no_label_map"})
        try:
            lb = base64.b64decode(label_map_b64)
            lab = Image.open(io.BytesIO(lb)).convert("L")
            label_map = np.array(lab, dtype=np.uint8)
        except Exception as e:
            return image_rgb, json.dumps({"error": "bad_label_map", "detail": str(e)})

        idx = int(float(mask_index)) if mask_index is not None else 1
        if idx <= 0:
            return image_rgb, json.dumps({"error": "invalid_mask_index"})

        mask = (label_map == idx)
        if not mask.any():
            return image_rgb, json.dumps({"error": "mask_not_found", "index": idx})

        # Limpieza basica
        mask = limpiar_mascara(mask)

        # Intentar importar el pipeline del Generador de Texturas
        gen_root_candidates = [
            Path("c:/Users/alane/OneDrive/Escritorio/Generandor de texturas"),
            Path("../Generandor de texturas"),
            Path("../../Generandor de texturas"),
            Path("Generandor de texturas"),
        ]
        gen_root = None
        for p in gen_root_candidates:
            if p.exists():
                gen_root = p.resolve()
                break
        if gen_root is not None:
            sys.path.insert(0, str(gen_root))
            sys.path.insert(0, str((gen_root / "pipeline")))

        try:
            from pipeline.perspective import PerspectiveAnalyzer
            from pipeline.depth_estimator import DepthEstimator
            from pipeline.renderer import TextureRenderer
        except Exception as e:
            return image_rgb, json.dumps({"error": "import_failed", "detail": str(e)})

        # Esquinas y perspectiva
        analyzer = PerspectiveAnalyzer()
        corners = analyzer.get_wall_corners(mask)

        # ── Validación de quad: si la máscara es dispersa (multi-máscara) las esquinas
        # pueden ser degeneradas (área casi 0). Fallback: bounding box de la máscara.
        def _quad_area(pts: np.ndarray) -> float:
            """Área del polígono usando fórmula de Shoelace."""
            n = len(pts)
            s = 0.0
            for _i in range(n):
                _j = (_i + 1) % n
                s += pts[_i, 0] * pts[_j, 1]
                s -= pts[_j, 0] * pts[_i, 1]
            return abs(s) / 2.0

        _MIN_QUAD_AREA_PX = 400  # < 20x20 px → degenerado
        if _quad_area(corners) < _MIN_QUAD_AREA_PX:
            _ys_m, _xs_m = np.where(mask)
            if len(_ys_m) > 0:
                # Usar bounding box de todos los píxeles de la máscara
                corners = np.array([
                    [float(_xs_m.min()), float(_ys_m.min())],  # TL
                    [float(_xs_m.max()), float(_ys_m.min())],  # TR
                    [float(_xs_m.max()), float(_ys_m.max())],  # BR
                    [float(_xs_m.min()), float(_ys_m.max())],  # BL
                ], dtype=float)

        # Profundidad y tamaño de pared
        depth_est = DepthEstimator(model_type="DPT_Large")
        depth_map = depth_est.estimate(image_rgb)
        wall_w_m, wall_h_m = depth_est.estimate_wall_size(mask, depth_map, corners)

        # Cargar textura (b64 o por nombre en textures)
        texture_np = None
        if texture_b64:
            try:
                tb = base64.b64decode(texture_b64)
                tpil = Image.open(io.BytesIO(tb)).convert("RGB")
                texture_np = np.array(tpil)
            except Exception:
                texture_np = None

        if texture_np is None and texture_name:
            tex_dir = gen_root / "textures" if gen_root is not None else Path("textures")
            cand = tex_dir / texture_name
            if cand.exists():
                texture_np = np.array(Image.open(cand).convert("RGB"))
            else:
                for ext in [".jpg", ".jpeg", ".png"]:
                    cc = tex_dir / (texture_name + ext)
                    if cc.exists():
                        texture_np = np.array(Image.open(cc).convert("RGB"))
                        break

        # Fallback texture
        if texture_np is None:
            texture_np = np.ones((256, 256, 3), dtype=np.uint8) * 200

        # Parsear params
        try:
            params = json.loads(params_json) if params_json else {}
        except Exception:
            params = {}

        tile_w_m = params.get("ancho_panel_m") or params.get("tile_w_m") or None
        tile_h_m = params.get("alto_panel_m") or params.get("tile_h_m") or None
        blend = float(params.get("intensidad_textura", params.get("blend_strength", 0.85)))
        separacion = int(params.get("separacion_vertical_px", params.get("separacion_px", 0)))
        separacion_h = int(params.get("separacion_horizontal_px", params.get("separacion_h_px", 0)))
        orientacion = params.get("orientacion", params.get("orientation", "vertical"))
        perspectiva_h = float(params.get("perspectiva_horizontal", params.get("perspectiva_h", 0.35)))
        perspectiva_v = float(params.get("perspectiva_vertical", params.get("perspectiva_v", 0.7)))
        modo_fusion = params.get("modo_fusion", params.get("blend_mode", "Luz suave"))

        # Ajuste de perspectiva según tipo de superficie (piso vs pared).
        # La orientación la controla exclusivamente el preset — no se sobreescribe aquí.
        surface_type = params.get("surface_type", "wall")
        if surface_type in ("floor", "deck", "ceiling"):
            # Para piso: perspectiva diferente a la de pared.
            # Solo se aplica si el preset NO tiene un valor explícito para estos campos.
            if "perspectiva_horizontal" not in params and "perspectiva_h" not in params:
                perspectiva_h = 0.15
            if "perspectiva_vertical" not in params and "perspectiva_v" not in params:
                perspectiva_v = 0.85

        # Ajuste orientacion
        if str(orientacion).lower().startswith("h"):
            texture_np = np.rot90(texture_np, k=1).copy()
            effective_tile_w = tile_h_m
            effective_tile_h = tile_w_m
        else:
            effective_tile_w = tile_w_m
            effective_tile_h = tile_h_m

        # Auto-remap de separación según orientación.
        # El usuario define UN valor de separación; el eje (v o h) depende de la orientación:
        #   - Modo horizontal (tablones izq-der): gaps entre tablones = h_bands (separacion_h_px)
        #   - Modo vertical   (tablones arriba-abajo): gaps entre tablones = v_bands (separacion_px)
        # Si el usuario sólo tiene un valor en el eje "opuesto", lo remapeamos automáticamente.
        _is_horiz = str(orientacion).lower().startswith("h")
        if _is_horiz and separacion_h == 0 and separacion > 0:
            # Vertical → Horizontal: mover sep_v a sep_h para que queden bandas h entre tablones
            separacion_h, separacion = separacion, 0
        elif not _is_horiz and separacion == 0 and separacion_h > 0:
            # Horizontal → Vertical: mover sep_h a sep_v para que queden bandas v entre tablones
            separacion, separacion_h = separacion_h, 0

        # Render — envuelto en try para capturar errores OpenCV residuales
        renderer = TextureRenderer()
        try:
            result = renderer.render(
                image=image_rgb,
                mask=mask,
                texture=texture_np,
                corners=corners,
                wall_w_m=wall_w_m,
                wall_h_m=wall_h_m,
                tile_w_m=effective_tile_w,
                tile_h_m=effective_tile_h,
                blend_strength=blend,
                separacion_px=separacion,
                separacion_h_px=separacion_h,
                horizontal_sep=_is_horiz,
                perspectiva_h=perspectiva_h,
                perspectiva_v=perspectiva_v,
                modo_fusion=modo_fusion,
            )
        except Exception as _render_err:
            print(f"  [Space render] Error en renderer.render: {_render_err}. Usando tiling simple.")
            # Fallback: tiling simple sobre la máscara sin perspectiva
            th_np, tw_np = texture_np.shape[:2]
            h_im, w_im = image_rgb.shape[:2]
            reps_y = max(1, -(-h_im // th_np))  # ceil division
            reps_x = max(1, -(-w_im // tw_np))
            big = np.tile(texture_np, (reps_y, reps_x, 1))[:h_im, :w_im]
            mask_3 = np.stack([mask, mask, mask], axis=2).astype(np.float32)
            orig_f = image_rgb.astype(np.float32) / 255.0
            big_f  = big.astype(np.float32) / 255.0
            result = np.clip(
                np.where(mask_3, big_f * blend + orig_f * (1.0 - blend), orig_f) * 255, 0, 255
            ).astype(np.uint8)

        combined = {
            "mask_index": idx,
            "wall_w_m": wall_w_m,
            "wall_h_m": wall_h_m,
            "params": params,
        }
        return result, 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})


# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(title="SAM Auto - Segmentacion") as demo:
    gr.Markdown("# Segmentacion Automatica - SAM ViT-Huge")
    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=16, maximum=128, step=16, value=PARAMS["points_per_batch"],
                label="points_per_batch  (no afecta calidad, solo velocidad)"
            )
            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",
    )

    # Endpoint oculto para renderizado remoto (usado por backend)
    _r_in = gr.Image(type="numpy", label="render_input", visible=False)
    _r_label = gr.Textbox(label="render_label_map_b64", visible=False)
    _r_mask = gr.Number(value=1, visible=False)
    _r_texture_name = gr.Textbox(label="texture_name", visible=False)
    _r_texture_b64 = gr.Textbox(label="texture_b64", visible=False)
    _r_params = gr.Textbox(label="render_params_json", visible=False)
    _r_out = gr.Image(type="numpy", label="render_output", visible=False)
    _r_json = gr.Textbox(label="render_json", visible=False)
    _r_btn = gr.Button(visible=False)
    _r_btn.click(
        fn=render_for_backend,
        inputs=[_r_in, _r_label, _r_mask, _r_texture_name, _r_texture_b64, _r_params],
        outputs=[_r_out, _r_json],
        api_name="render",
    )

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