Upload app.py

app.py

@@ -10,11 +10,30 @@ 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)
-
+# ── 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,
+    "crops_n_layers":         1,
+}
+
 
 # ── Renderizado ───────────────────────────────────────────────────────────────
 def _render_masks(imagen_rgb: Image.Image, masks: list) -> Image.Image:
@@ -28,28 +47,53 @@ def _render_masks(imagen_rgb: Image.Image, masks: list) -> Image.Image:
     return Image.fromarray(blended)
 
 
-# ── Segmentación UI ───────────────────────────────────────────────────────────
-def segmentar(imagen: Image.Image):
+def _load_pipeline():
     global sam_vit_pipeline
-    if imagen is None:
-        return None, "Sube una imagen para comenzar."
-
     if sam_vit_pipeline is None:
-        print("Cargando SAM ViT-Huge (CPU)…")
+        print("Cargando SAM ViT-Huge...")
         sam_vit_pipeline = hf_pipeline(
             "mask-generation",
             model="facebook/sam-vit-huge",
-            device=-1,
+            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,
+    crops_n_layers: int,
+):
+    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),
+        "crops_n_layers":         int(crops_n_layers),
+    })
+
+    _load_pipeline()
     imagen_rgb = imagen.convert("RGB")
     resultado = sam_vit_pipeline(
         imagen_rgb,
-        points_per_batch=32,
-        pred_iou_thresh=0.95,
-        stability_score_thresh=0.5,
-        min_mask_region_area=4500,
-        box_nms_thresh=0.8,
+        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"],
+        crops_n_layers=PARAMS["crops_n_layers"],
     )
     if isinstance(resultado, list):
         resultado = resultado[0]
@@ -58,14 +102,23 @@ def segmentar(imagen: Image.Image):
     if not masks:
         return imagen_rgb, "No se detectaron zonas."
 
-    return _render_masks(imagen_rgb, masks), f"✅ {len(masks)} zonas detectadas automáticamente"
+    info = (
+        f"UI: {len(masks)} zonas  |  "
+        f"iou={PARAMS['pred_iou_thresh']}  stab={PARAMS['stability_score_thresh']}  "
+        f"crops={PARAMS['crops_n_layers']}  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 vía gradio_client (api_name='/segment').
-    Entrada : numpy uint8 H×W×3.
+    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:
@@ -73,25 +126,18 @@ def segment_for_backend(image_np: np.ndarray):
             empty = np.zeros((100, 100, 3), dtype=np.uint8)
             return empty, json.dumps({"masks": [], "label_map_b64": ""})
 
-        global sam_vit_pipeline
-        if sam_vit_pipeline is None:
-            print("segment_for_backend: cargando SAM ViT-Huge…")
-            sam_vit_pipeline = hf_pipeline(
-                "mask-generation",
-                model="facebook/sam-vit-huge",
-                device=-1,
-            )
-
+        _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=32,
-            pred_iou_thresh=0.95,
-            stability_score_thresh=0.5,
-            min_mask_region_area=4500,
-            box_nms_thresh=0.8,
+            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"],
+            crops_n_layers=PARAMS["crops_n_layers"],
         )
         if isinstance(resultado, list):
             resultado = resultado[0]
@@ -99,7 +145,7 @@ def segment_for_backend(image_np: np.ndarray):
         all_masks_raw = resultado.get("masks", [])
         masks_bool = [np.array(m).astype(bool) for m in all_masks_raw]
 
-        # Label map: cada píxel contiene el índice de la máscara (1-based)
+        # 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):
@@ -117,21 +163,20 @@ def segment_for_backend(image_np: np.ndarray):
                 "bbox_xywh": bbox,
             })
 
-        # Codificar label map como PNG en base64
         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 de visualización
         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": "auto",
-            "motor": "SAM Auto (Exteriores - Sin Prompts)",
+            "entorno": "gpu",
+            "motor": "SAM Auto (GPU - ZeroGPU)",
+            "params_used": dict(PARAMS),
         }
         return overlay_np, json.dumps(combined, ensure_ascii=False)
 
@@ -142,39 +187,68 @@ def segment_for_backend(image_np: np.ndarray):
 
 
 # ── UI ────────────────────────────────────────────────────────────────────────
-def crear_app():
-    with gr.Blocks(title="SAM Auto — CPU") as demo:
-        gr.Markdown("# 🏗️ Segmentación Automática — SAM ViT-Huge (CPU)")
+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(
-            "SAM detecta todos los elementos de la imagen de forma automática, "
-            "sin necesidad de seleccionar zonas ni escribir prompts."
+            "> 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():
-            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")
-
-        boton.click(fn=segmentar, inputs=imagen_entrada, outputs=[imagen_salida, estado])
-        imagen_entrada.upload(fn=segmentar, inputs=imagen_entrada, outputs=[imagen_salida, estado])
-
-        # Endpoint oculto para el backend Docker — componentes en nivel raíz del Blocks
-        _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
+            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)"
+            )
+            sl_crops = gr.Slider(
+                minimum=0, maximum=2, step=1, value=PARAMS["crops_n_layers"],
+                label="crops_n_layers  (0=imagen completa · 1=+4 recortes · 2=+16 recortes, lento)"
+            )
 
+    all_inputs = [imagen_entrada, sl_pred_iou, sl_stability, sl_batch, sl_min_area, sl_nms, sl_crops]
+    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",
+    )
 
-demo = crear_app()
 if __name__ == "__main__":
     demo.launch()