Upload app.py

app.py

@@ -10,28 +10,18 @@ import cv2
 from PIL import Image
 from transformers import pipeline as hf_pipeline
 
-# ── 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
+# CPU-only: use more threads for better throughput
+torch.set_num_threads(4)
+
 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,
+    "points_per_batch":       16,
     "min_mask_region_area":   4500,
     "box_nms_thresh":         0.8,
-    "crops_n_layers":         1,
 }
 
 
@@ -50,17 +40,15 @@ def _render_masks(imagen_rgb: Image.Image, masks: list) -> Image.Image:
 def _load_pipeline():
     global sam_vit_pipeline
     if sam_vit_pipeline is None:
-        print("Cargando SAM ViT-Huge...")
+        print("Cargando SAM ViT-Huge (CPU)...")
         sam_vit_pipeline = hf_pipeline(
             "mask-generation",
             model="facebook/sam-vit-huge",
-            device=DEVICE,
+            device=-1,
         )
 
 
 # ── Segmentacion UI ───────────────────────────────────────────────────────────
-@spaces.GPU
-@torch.no_grad()
 def segmentar(
     imagen: Image.Image,
     pred_iou_thresh: float,
@@ -68,20 +56,17 @@ def segmentar(
     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()
@@ -93,7 +78,6 @@ def segmentar(
         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]
@@ -105,15 +89,13 @@ def segmentar(
     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"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').
@@ -137,7 +119,6 @@ def segment_for_backend(image_np: np.ndarray):
             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]
@@ -145,7 +126,6 @@ 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 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):
@@ -174,8 +154,8 @@ def segment_for_backend(image_np: np.ndarray):
         combined = {
             "masks": masks_out,
             "label_map_b64": label_map_b64,
-            "entorno": "gpu",
-            "motor": "SAM Auto (GPU - ZeroGPU)",
+            "entorno": "cpu",
+            "motor": "SAM Auto (CPU)",
             "params_used": dict(PARAMS),
         }
         return overlay_np, json.dumps(combined, ensure_ascii=False)
@@ -187,68 +167,69 @@ def segment_for_backend(image_np: np.ndarray):
 
 
 # ── 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):
+def crear_app():
+    with gr.Blocks(title="SAM Auto - CPU") as demo:
+        gr.Markdown("# Segmentacion Automatica - SAM ViT-Huge (CPU)")
         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."
+            "SAM detecta todos los elementos de la imagen de forma automatica, "
+            "sin necesidad de seleccionar zonas ni escribir prompts."
         )
+
         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)"
-            )
-            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)"
+            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
 
-    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()