Update app.py

app.py

@@ -1,3 +1,4 @@
+#region Imports
 import os
 
 # Route caches to /tmp to avoid filling the Space persistent storage
@@ -22,24 +23,20 @@ try:
     _ULTRA_OK = True
 except Exception:
     _ULTRA_OK = False
+#endregion
 
-# Config
-MAX_SIDE_PX = 70  # filtro lato massimo bbox per modello A (SAHI)
+#region Config and setup
+MAX_SIDE_PX = -1 # 70  # possible filtering for max bbox side for Model A
 SEG_DEFAULT_ALPHA = 0.45
-
+# Weights
+WEIGHTS_DETECTION = "weights/berries.pt"
+WEIGHTS_SEGMENTATION = "weights/bunches.pt"
 # Simple global caches to avoid reloading models each click
 _DET_MODEL_CACHE = {}  # key: (weights_path, device) -> AutoDetectionModel
 _SEG_MODEL_CACHE = {}  # key: weights_path -> YOLO
+#endregion
 
-def _ensure_rgb(img: np.ndarray) -> np.ndarray:
-    if img is None:
-        return None
-    if img.ndim == 2:
-        return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
-    if img.shape[2] == 4:
-        return cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
-    return img
-
+#region Model and device handling
 def _choose_device(user_choice: str) -> str:
     if user_choice != "auto":
         return user_choice
@@ -54,14 +51,13 @@ def _get_det_model(weights_path: str, device: str, conf: float):
     Returns a cached SAHI AutoDetectionModel. Updates confidence on the fly.
     """
     if not os.path.exists(weights_path):
-        raise gr.Error(f"Pesi detection non trovati: {weights_path}")
+        raise gr.Error(f"Detection weights not found: {weights_path}")
     key = (weights_path, device)
     model = _DET_MODEL_CACHE.get(key)
     if model is None:
-        # SAHI uses yolov8 wrapper for Ultralytics models (works for v8/v9/v11)
         try:
             model = AutoDetectionModel.from_pretrained(
-                model_type="yolov8",
+                model_type="yolov11",
                 model_path=weights_path,
                 confidence_threshold=conf,
                 device=device,
@@ -69,7 +65,7 @@ def _get_det_model(weights_path: str, device: str, conf: float):
         except Exception:
             # CPU fallback
             model = AutoDetectionModel.from_pretrained(
-                model_type="yolov8",
+                model_type="yolov11",
                 model_path=weights_path,
                 confidence_threshold=conf,
                 device="cpu",
@@ -85,29 +81,96 @@ def _get_det_model(weights_path: str, device: str, conf: float):
 
 def _get_seg_model(weights_path: str):
     if not _ULTRA_OK:
-        raise gr.Error("Ultralytics non installato. Installa con: pip install ultralytics")
+        raise gr.Error("Ultralytics not found, please install it with: pip install ultralytics")
     if not os.path.exists(weights_path):
-        raise gr.Error(f"Pesi segmentation non trovati: {weights_path}")
+        raise gr.Error(f"Segmentation weights not found: {weights_path}")
     model = _SEG_MODEL_CACHE.get(weights_path)
     if model is None:
         model = YOLO(weights_path)
         _SEG_MODEL_CACHE[weights_path] = model
     return model
+#endregion
+
+#region Inference
+def _sahi_predict(image_rgb: np.ndarray, det_model, slice_h, slice_w, overlap_h, overlap_w):
+    return get_sliced_prediction(
+        image_rgb,
+        det_model,
+        slice_height=int(slice_h),
+        slice_width=int(slice_w),
+        overlap_height_ratio=float(overlap_h),
+        overlap_width_ratio=float(overlap_w),
+        postprocess_class_agnostic=False,
+        verbose=0,
+    )
+
+def run_det(
+    image, state,
+    weights_det_path, conf_det, slice_h, slice_w, overlap_h, overlap_w, device,
+    target_class, use_target, auto_opt_slice
+):
+    """
+    Run model A for berries and updates only 'det' overlay.
+    Assemble final image with both layers (det + seg) in the proper order.
+    """
+    if state is None or state.get("base") is None:
+        raise gr.Error("Loading an image is required before inference.")
+    base = state["base"]
+    det_model = _get_det_model(weights_det_path, _choose_device(device), conf_det)
+    sahi_res = _sahi_predict(base, det_model, slice_h, slice_w, overlap_h, overlap_w)
+
+    overlay_rgb, alpha, counts = _draw_boxes_overlay(base, sahi_res, target_class, bool(use_target))
+
+    state["det"] = {"overlay": overlay_rgb, "alpha": alpha, "ts": time.time()}
+    state["det_counts"] = counts
+
+    layers = [state["det"], state.get("seg")]
+    composite = _composite_layers(base, layers)
+    return composite, state, state["det_counts"], state.get("seg_counts", "")
+
+def run_seg(
+    image, state,
+    weights_seg_path, conf_seg, device,
+    target_class, use_target, seg_alpha
+):
+    """
+    Run model B for bunches and updates only 'seg' overlay.
+    Assemble final image with both layers (det + seg) in the proper order.
+    """
+    if state is None or state.get("base") is None:
+        raise gr.Error("Loading an image is required before inference.")
+    base = state["base"]
+    seg_model = _get_seg_model(weights_seg_path)
+    # device is handled in predict
+    try:
+        seg_results = seg_model.predict(source=base, conf=float(conf_seg), device=_choose_device(device), verbose=False)
+        r0 = seg_results[0] if isinstance(seg_results, (list, tuple)) else seg_results
+    except Exception as e:
+        raise gr.Error(f"Error in segmentation inference: {e}")
 
-def _optimize_slicing_dims(H: int, W: int, slice_h: int, slice_w: int, overlap_h: float, overlap_w: float, auto_opt: bool):
-    if not auto_opt:
-        return int(slice_h), int(slice_w), float(overlap_h), float(overlap_w)
-    sh = min(int(slice_h), H)
-    sw = min(int(slice_w), W)
-    # If the image already fits in one slice, remove overlap to reduce work
-    oh = 0.0 if (H <= sh and W <= sw) else float(overlap_h)
-    ow = 0.0 if (H <= sh and W <= sw) else float(overlap_w)
-    return sh, sw, oh, ow
+    overlay_rgb, alpha, counts = _draw_seg_overlay(base, r0, target_class, bool(use_target), float(seg_alpha))
+    state["seg"] = {"overlay": overlay_rgb, "alpha": alpha, "ts": time.time()}
+    state["seg_counts"] = counts
+
+    layers = [state.get("det"), state["seg"]]
+    composite = _composite_layers(base, layers)
+    return composite, state, state.get("det_counts", ""), state["seg_counts"]
+#endregion
+
+#region Draw
+def _ensure_rgb(img: np.ndarray) -> np.ndarray:
+    if img is None:
+        return None
+    if img.ndim == 2:
+        return cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    if img.shape[2] == 4:
+        return cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
+    return img
 
 def _draw_boxes_overlay(image_rgb: np.ndarray, sahi_result, target_class: str, use_target: bool):
     """
     Returns overlay_rgb (H,W,3), alpha_mask (H,W) uint8, counts_text
-    Only draws rectangles (no labels). Filters boxes with max side > MAX_SIDE_PX.
+    Only draws rectangles (no labels). Filters boxes with max side > MAX_SIDE_PX if MAX_SIDE_PX > 0 (hard coded).
     """
     H, W = image_rgb.shape[:2]
     overlay = np.zeros((H, W, 3), dtype=np.uint8)
@@ -135,8 +198,9 @@ def _draw_boxes_overlay(image_rgb: np.ndarray, sahi_result, target_class: str, u
         h = max(0, y2 - y1)
         if w == 0 or h == 0:
             continue
-        if max(w, h) > MAX_SIDE_PX:
-            continue
+        if MAX_SIDE_PX > 0:
+            if max(w, h) > MAX_SIDE_PX:
+                continue
 
         area = getattr(item.bbox, "area", w * h)
         try:
@@ -163,9 +227,9 @@ def _draw_boxes_overlay(image_rgb: np.ndarray, sahi_result, target_class: str, u
     # Convert overlay BGR -> RGB
     overlay_rgb = cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB)
     if use_target:
-        counts = f"target='{target_class}': {target_count} | totale: {total_count}"
+        counts = f"target='{target_class}': {target_count} | total: {total_count}"
     else:
-        counts = f"totale: {total_count}"
+        counts = f"total: {total_count}"
     return overlay_rgb, alpha, counts
 
 def _draw_seg_overlay(image_rgb: np.ndarray, yolo_result, target_class: str, use_target: bool, fill_alpha: float = SEG_DEFAULT_ALPHA):
@@ -184,7 +248,7 @@ def _draw_seg_overlay(image_rgb: np.ndarray, yolo_result, target_class: str, use
     masks = getattr(r, "masks", None)
 
     if boxes is None or len(boxes) == 0:
-        counts = f"target='{target_class}': 0 | totale: 0" if use_target else "totale: 0"
+        counts = f"target='{target_class}': 0 | total: 0" if use_target else "total: 0"
         return cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB), alpha, counts
 
     N = len(boxes)
@@ -255,9 +319,9 @@ def _draw_seg_overlay(image_rgb: np.ndarray, yolo_result, target_class: str, use
 
     overlay_rgb = cv2.cvtColor(overlay_bgr, cv2.COLOR_BGR2RGB)
     if use_target:
-        counts = f"target='{target_class}': {target_count} | totale: {total_count}"
+        counts = f"target='{target_class}': {target_count} | total: {total_count}"
     else:
-        counts = f"totale: {total_count}"
+        counts = f"total: {total_count}"
     return overlay_rgb, alpha, counts
 
 def _composite_layers(base_rgb: np.ndarray, layers: list):
@@ -285,23 +349,9 @@ def _composite_layers(base_rgb: np.ndarray, layers: list):
 
     return np.clip(result, 0, 255).astype(np.uint8)
 
-def _sahi_predict(image_rgb: np.ndarray, det_model, slice_h, slice_w, overlap_h, overlap_w):
-    return get_sliced_prediction(
-        image_rgb,
-        det_model,
-        slice_height=int(slice_h),
-        slice_width=int(slice_w),
-        overlap_height_ratio=float(overlap_h),
-        overlap_width_ratio=float(overlap_w),
-        postprocess_class_agnostic=False,
-        verbose=0,
-    )
-
-# Gradio callables
-
 def on_image_upload(image, state):
     """
-    Resetta gli overlay quando si carica una nuova immagine.
+    Reset overlays if uploading a new image.
     """
     if image is None:
         return None, {"base": None, "det": None, "seg": None, "det_counts": "", "seg_counts": ""}, "", ""
@@ -309,60 +359,6 @@ def on_image_upload(image, state):
     new_state = {"base": img_rgb, "det": None, "seg": None, "det_counts": "", "seg_counts": ""}
     return img_rgb, new_state, "", ""
 
-def run_det(
-    image, state,
-    weights_det_path, conf_det, slice_h, slice_w, overlap_h, overlap_w, device,
-    target_class, use_target, auto_opt_slice
-):
-    """
-    Esegue il modello A (SAHI detection) e aggiorna solo l'overlay 'det'.
-    Recompone l'immagine finale con entrambi i layer (det + seg) nell'ordine temporale.
-    """
-    if state is None or state.get("base") is None:
-        raise gr.Error("Carica prima un'immagine.")
-    base = state["base"]
-    H, W = base.shape[:2]
-    det_model = _get_det_model(weights_det_path, _choose_device(device), conf_det)
-    sh, sw, oh, ow = _optimize_slicing_dims(H, W, slice_h, slice_w, overlap_h, overlap_w, auto_opt_slice)
-    sahi_res = _sahi_predict(base, det_model, sh, sw, oh, ow)
-
-    overlay_rgb, alpha, counts = _draw_boxes_overlay(base, sahi_res, target_class, bool(use_target))
-
-    state["det"] = {"overlay": overlay_rgb, "alpha": alpha, "ts": time.time()}
-    state["det_counts"] = counts
-
-    layers = [state["det"], state.get("seg")]
-    composite = _composite_layers(base, layers)
-    return composite, state, state["det_counts"], state.get("seg_counts", "")
-
-def run_seg(
-    image, state,
-    weights_seg_path, conf_seg, device,
-    target_class, use_target, seg_alpha
-):
-    """
-    Esegue il modello B (YOLO segmentation) e aggiorna solo l'overlay 'seg'.
-    Recompone l'immagine finale con entrambi i layer (det + seg) nell'ordine temporale.
-    """
-    if state is None or state.get("base") is None:
-        raise gr.Error("Carica prima un'immagine.")
-    base = state["base"]
-    seg_model = _get_seg_model(weights_seg_path)
-    # device is handled in predict
-    try:
-        seg_results = seg_model.predict(source=base, conf=float(conf_seg), device=_choose_device(device), verbose=False)
-        r0 = seg_results[0] if isinstance(seg_results, (list, tuple)) else seg_results
-    except Exception as e:
-        raise gr.Error(f"Errore inferenza segmentation: {e}")
-
-    overlay_rgb, alpha, counts = _draw_seg_overlay(base, r0, target_class, bool(use_target), float(seg_alpha))
-    state["seg"] = {"overlay": overlay_rgb, "alpha": alpha, "ts": time.time()}
-    state["seg_counts"] = counts
-
-    layers = [state.get("det"), state["seg"]]
-    composite = _composite_layers(base, layers)
-    return composite, state, state.get("det_counts", ""), state["seg_counts"]
-
 def clear_overlays(image, state):
     if state is None or state.get("base") is None:
         return None, {"base": None, "det": None, "seg": None, "det_counts": "", "seg_counts": ""}, "", ""
@@ -372,9 +368,9 @@ def clear_overlays(image, state):
     state["det_counts"] = ""
     state["seg_counts"] = ""
     return base, state, "", ""
+#endregion
 
-# Maintenance helpers
-
+#region Maintenance
 def _dir_size(path: str) -> int:
     try:
         total = 0
@@ -437,37 +433,24 @@ def clean_caches():
         except Exception:
             pass
     return "Removed:\n" + ("\n".join(removed) if removed else "(none)")
+#endregion
 
 def build_app():
-    with gr.Blocks(title="YOLOv11 SAHI Detection + YOLO Segmentation (dual overlays)") as demo:
+    with gr.Blocks(title="Berries detection & bunches segmentation") as demo:
         gr.Markdown(
-            "## Doppia inferenza su stessa immagine, overlay combinati\n"
-            "- Modello A: SAHI detection (usa pesi YOLOv11 seg come detection) — solo bbox, filtro lato > 70px.\n"
-            "- Modello B: YOLO segmentation nativo — maschere riempite + contorno.\n"
-            "- Esegui i modelli con pulsanti separati; gli overlay si accumulano sull'immagine base (nuovo overlay sopra).\n"
-            "- Opzionale: disabilita l'evidenziazione della classe target se non ti serve."
+            "## Double inference on the same image with combined overlays\n"
+            "- Model A: berries detection\n"
+            "- Model B: bunches segmentation\n"
+            "- Individual executions\n"
         )
 
         state = gr.State({"base": None, "det": None, "seg": None, "det_counts": "", "seg_counts": ""})
 
         with gr.Row():
             with gr.Column(scale=1):
-                img_in = gr.Image(label="Immagine", type="numpy")
-                with gr.Accordion("Pesi modelli", open=True):
-                    weights_det = gr.Textbox(
-                        label="Pesi Modello A (Detection + SAHI, .pt)",
-                        value="weights/best.pt",
-                    )
-                    weights_seg = gr.Textbox(
-                        label="Pesi Modello B (Segmentation, .pt)",
-                        value="weights/seg.pt",
-                    )
-
-                with gr.Row():
-                    target = gr.Textbox(label="Classe target", value="berry")
-                    use_target = gr.Checkbox(label="Usa classe target", value=True)
+                img_in = gr.Image(label="Image", type="numpy")
 
-                with gr.Tab("Modello A — SAHI Detection"):
+                with gr.Tab("Model A — Berries Detection"):
                     with gr.Row():
                         conf_det = gr.Slider(0.0, 1.0, value=0.35, step=0.01, label="Confidence (A)")
                         device_a = gr.Dropdown(["auto", "cuda:0", "cpu"], value="auto", label="Device")
@@ -477,29 +460,28 @@ def build_app():
                     with gr.Row():
                         overlap_h = gr.Slider(0.0, 0.9, value=0.10, step=0.01, label="Overlap H (A)")
                         overlap_w = gr.Slider(0.0, 0.9, value=0.10, step=0.01, label="Overlap W (A)")
-                    auto_opt_slice = gr.Checkbox(label="Ottimizza slicing su immagini piccole", value=True)
-                    btn_det = gr.Button("Esegui Modello A (SAHI)")
+                    btn_det = gr.Button("Run berries detection")
 
-                with gr.Tab("Modello B — YOLO Segmentation"):
+                with gr.Tab("Model B — Bunches Segmentation"):
                     with gr.Row():
                         conf_seg = gr.Slider(0.0, 1.0, value=0.35, step=0.01, label="Confidence (B)")
-                        seg_alpha = gr.Slider(0.0, 1.0, value=SEG_DEFAULT_ALPHA, step=0.05, label="Alpha maschere (B)")
+                        seg_alpha = gr.Slider(0.0, 1.0, value=SEG_DEFAULT_ALPHA, step=0.05, label="Alpha masks (B)")
                         device_b = gr.Dropdown(["auto", "cuda:0", "cpu"], value="auto", label="Device")
-                    btn_seg = gr.Button("Esegui Modello B (Seg)")
+                    btn_seg = gr.Button("Run bunches segmentation")
 
                 with gr.Row():
-                    btn_clear = gr.Button("Pulisci overlay", variant="secondary")
+                    btn_clear = gr.Button("Clean overlay", variant="secondary")
 
-                with gr.Accordion("Manutenzione spazio", open=False):
-                    btn_check = gr.Button("Controlla storage")
-                    btn_clean = gr.Button("Pulisci cache")
-                    maint_out = gr.Textbox(label="Log manutenzione", interactive=False)
+                with gr.Accordion("Disk Maintenance", open=False):
+                    btn_check = gr.Button("Check storage")
+                    btn_clean = gr.Button("Clean cache")
+                    maint_out = gr.Textbox(label="Log Maintenance", interactive=False)
 
             with gr.Column(scale=2):
-                img_out = gr.Image(label="Risultato combinato", type="numpy")
+                img_out = gr.Image(label="Combined Result", type="numpy")
                 with gr.Row():
-                    counts_out_det = gr.Textbox(label="Conteggi (A)", interactive=False)
-                    counts_out_seg = gr.Textbox(label="Conteggi (B)", interactive=False)
+                    counts_out_det = gr.Textbox(label="Counts - Berries", interactive=False)
+                    counts_out_seg = gr.Textbox(label="Counts - Bunches", interactive=False)
 
         # Wiring
         img_in.change(
@@ -512,8 +494,8 @@ def build_app():
             run_det,
             inputs=[
                 img_in, state,
-                weights_det, conf_det, slice_h, slice_w, overlap_h, overlap_w, device_a,
-                target, use_target, auto_opt_slice
+                WEIGHTS_DETECTION, conf_det, slice_h, slice_w, overlap_h, overlap_w, device_a#,
+                #target, use_target, auto_opt_slice
             ],
             outputs=[img_out, state, counts_out_det, counts_out_seg],
         )
@@ -522,8 +504,9 @@ def build_app():
             run_seg,
             inputs=[
                 img_in, state,
-                weights_seg, conf_seg, device_b,
-                target, use_target, seg_alpha
+                WEIGHTS_SEGMENTATION, conf_seg, device_b,
+                seg_alpha
+                #target, use_target, seg_alpha
             ],
             outputs=[img_out, state, counts_out_det, counts_out_seg],
         )