Swap classical region-grow for nnInteractive in segment_from_prompt seam

flask-server/api/api_blueprint.py

@@ -5143,7 +5143,8 @@ def interactive_segment(case_id):
         # float32: half the RAM of nibabel's float64 default — these are public
         # endpoints and a full-res CT at float64 is multiple GB per request.
         ct = ct_obj.get_fdata(dtype=np.float32)
-        mask = segment_from_prompt(ct, ct_obj.affine, body)
+        case_key = f"{case_id}:{'low' if low else 'full'}"
+        mask = segment_from_prompt(ct, ct_obj.affine, body, case_key=case_key)
         if int(mask.sum()) == 0:
             return jsonify({"error": "Nothing grew from that point — try a different spot or a higher tolerance."}), 422
 

flask-server/services/advanced_analysis.py

@@ -99,15 +99,14 @@ def region_grow(
     # Fill interior holes so the proposal is a solid object.
     mask = ndimage.binary_fill_holes(mask)
     return mask.astype(np.uint8)
+USE_NNINTERACTIVE = True
 
 
-def segment_from_prompt(ct: np.ndarray, affine: np.ndarray, prompt: dict) -> np.ndarray:
+def segment_from_prompt(ct: np.ndarray, affine: np.ndarray, prompt: dict, case_key: str | None = None) -> np.ndarray:
     """Model-agnostic entry point for the click-to-segment tool.
 
-    `prompt` carries a seed (`point_lps` or `point_ijk`), optional `tolerance`
-    and `box_lps`. This is the seam to swap in a promptable foundation model:
-    replace the region_grow call with `medsam2.infer(ct, prompt)` returning a
-    mask of the same shape — nothing upstream or in the frontend changes.
+    `case_key` (e.g. "17:full") lets the nnInteractive path cache the
+    uploaded volume across requests for the same case+resolution.
     """
     if "point_ijk" in prompt:
         seed = tuple(int(v) for v in prompt["point_ijk"])
@@ -124,15 +123,24 @@ def segment_from_prompt(ct: np.ndarray, affine: np.ndarray, prompt: dict) -> np.
             tuple(min(a, b) for a, b in zip(c0, c1)),
             tuple(max(a, b) + 1 for a, b in zip(c0, c1)),
         )
-    # The prompt comes straight from a public endpoint — clamp client numerics.
-    tolerance = min(max(float(prompt.get("tolerance", 80.0)), 1.0), 1000.0)
-    return region_grow(
-        ct,
-        seed,
-        tolerance=tolerance,
-        box_ijk=box_ijk,
-    )
 
+    if USE_NNINTERACTIVE:
+        try:
+            from services.nninteractive_predictor import predict
+            mask = predict(
+                ct,
+                case_key or "unkeyed",
+                point_ijk=seed if box_ijk is None else None,
+                box_ijk=box_ijk,
+            )
+            if mask.sum() > 0:
+                return mask
+            print("[segment_from_prompt] nnInteractive returned empty mask, falling back to region_grow")
+        except Exception as e:
+            print(f"[segment_from_prompt] nnInteractive failed ({type(e).__name__}: {e}), falling back to region_grow")
+
+    tolerance = min(max(float(prompt.get("tolerance", 80.0)), 1.0), 1000.0)
+    return region_grow(ct, seed, tolerance=tolerance, box_ijk=box_ijk)
 
 # --------------------------------------------------------------------------- #
 # 2. Vessel curved-planar analysis

flask-server/services/nninteractive_predictor.py

@@ -0,0 +1,100 @@
+"""
+flask-server/services/nninteractive_predictor.py
+
+Talks to the standalone `nninteractive-server` process on bdmap1
+(127.0.0.1:1527, model + GPU loaded once at server startup) via
+nnInteractiveRemoteInferenceSession.
+
+CONFIRMED end-to-end on bdmap1 against PanTS_00000001 (2026-08-15):
+  - add_point_interaction(coords, include_interaction=True)
+      coords = [i, j, k]  -> works, produced 142284 voxels on a test seed.
+  - add_bbox_interaction(bbox, include_interaction=True)
+      bbox = [[x_lo, x_hi], [y_lo, y_hi], [z_lo, z_hi]]  (per-axis pairs,
+      NOT two corner points). Exactly one axis must have size == 1 (a 2D
+      box on a single slice) -- size == 0 raises ValueError, and all three
+      axes > 1 raises "3D bounding box... not supported by the loaded
+      model checkpoint" (this checkpoint is 2D-box-only). Produced 16227
+      voxels on a 30x30 test box.
+
+Every box prompt is flattened to zero thickness on whichever axis has the
+smallest extent -- see _corners_to_axis_pairs(). This matches a box drawn
+on one 2D viewport pane, but has NOT yet been verified against a real
+frontend box-drag; verify this once wired up.
+
+Since api_blueprint.py's `_ANALYSIS_SLOTS` semaphore already serializes all
+calls into `interactive_segment()`, one shared session with no extra locking
+here is safe. Gunicorn is `--workers 1 --threads 8` (single process), so this
+module-level cache is correctly shared across every request thread.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+SERVER_URL = "http://127.0.0.1:1527"
+
+_session = None
+_cached_case_key: str | None = None
+_cached_ct_shape: tuple | None = None
+_target_buffer: np.ndarray | None = None
+
+
+def _get_session():
+    global _session
+    if _session is None:
+        from nnInteractive.inference.remote.remote_session import nnInteractiveRemoteInferenceSession
+        _session = nnInteractiveRemoteInferenceSession(server_url=SERVER_URL)
+        if not _session.ping():
+            raise RuntimeError(
+                f"nninteractive-server not reachable at {SERVER_URL} — "
+                "check it's running (tmux session 'nninteractive' on bdmap1)."
+            )
+    return _session
+
+
+def _ensure_volume_loaded(ct: np.ndarray, case_key: str) -> None:
+    global _cached_case_key, _cached_ct_shape, _target_buffer
+    session = _get_session()
+    if _cached_case_key == case_key and _cached_ct_shape == ct.shape:
+        return
+    session.set_image(ct[None])
+    _target_buffer = np.zeros(ct.shape, dtype=np.uint8)
+    session.set_target_buffer(_target_buffer)
+    _cached_case_key = case_key
+    _cached_ct_shape = ct.shape
+
+
+def _corners_to_axis_pairs(lo, hi) -> list[list[int]]:
+    lo, hi = list(lo), list(hi)
+    extents = [hi[d] - lo[d] for d in range(3)]
+    flatten_axis = min(range(3), key=lambda d: extents[d])
+    pairs = []
+    for d in range(3):
+        if d == flatten_axis:
+            start = lo[d]
+            pairs.append([start, start + 1])
+        else:
+            end = hi[d] if hi[d] > lo[d] else lo[d] + 1
+            pairs.append([lo[d], end])
+    return pairs
+
+
+def predict(
+    ct: np.ndarray,
+    case_key: str,
+    point_ijk=None,
+    box_ijk=None,
+) -> np.ndarray:
+    session = _get_session()
+    _ensure_volume_loaded(ct, case_key)
+    session.reset_interactions()
+
+    if point_ijk is not None:
+        session.add_point_interaction(list(point_ijk), include_interaction=True)
+    elif box_ijk is not None:
+        lo, hi = box_ijk
+        axis_pairs = _corners_to_axis_pairs(lo, hi)
+        session.add_bbox_interaction(axis_pairs, include_interaction=True)
+    else:
+        raise ValueError("predict() needs point_ijk or box_ijk")
+
+    return _target_buffer.copy()