Option B: dual-venv subprocess dispatch (venv_v3=Cellpose3.1, venv_v4=CellposeSAM)

app.py

@@ -6,6 +6,9 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 
+from setup_venvs import setup_venvs
+setup_venvs()  # create venv_v3 / venv_v4 on first launch
+
 from segmentation import MODELS, run_segmentation
 from mask_utils import (
     create_colored_overlay,

requirements.txt

@@ -1,8 +1,6 @@
-cellpose>=3.1.1
 gradio>=5.0.0
 numpy<2
-opencv-python-headless
 Pillow
 scikit-image
 pandas
-timm
+opencv-python-headless

requirements_v3.txt

@@ -0,0 +1,5 @@
+cellpose==3.1.1.2
+numpy<2
+opencv-python-headless
+Pillow
+scikit-image

requirements_v4.txt

@@ -0,0 +1,6 @@
+cellpose
+numpy<2
+opencv-python-headless
+Pillow
+scikit-image
+timm

segmentation.py

@@ -1,40 +1,21 @@
+import json
+import subprocess
+import tempfile
+from pathlib import Path
+
 import numpy as np
-import cv2
 from PIL import Image
-from cellpose import models
 
+from setup_venvs import get_python_executable
+
+WORKER = str(Path(__file__).parent / "worker.py")
 
 MODELS = {
     "CellposeSAM": ["cpsam"],
-    "Cellpose3.1": ["cyto3"],
+    "Cellpose3.1": ["cyto3", "nuclei"],
 }
 
 
-def load_image(image_input) -> np.ndarray:
-    """Accept a numpy array (from Gradio) or file path, return H×W×3 uint8 RGB."""
-    if isinstance(image_input, np.ndarray):
-        img = image_input.copy()
-        if img.ndim == 2:
-            img = np.stack([img, img, img], axis=-1)
-        elif img.shape[2] == 4:
-            img = img[:, :, :3]
-        return img.astype(np.uint8)
-
-    # file path
-    img = cv2.imread(str(image_input), cv2.IMREAD_UNCHANGED)
-    if img is None:
-        img = np.array(Image.open(str(image_input)).convert("RGB"))
-        return img.astype(np.uint8)
-
-    if img.ndim == 2:
-        img = np.stack([img, img, img], axis=-1)
-    elif img.shape[2] == 4:
-        img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
-    else:
-        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    return img.astype(np.uint8)
-
-
 def run_segmentation(
     image_input,
     model_type: str,
@@ -46,73 +27,58 @@ def run_segmentation(
     use_gpu: bool,
 ) -> tuple[np.ndarray, int]:
     """
-    Run Cellpose segmentation.
-
-    Returns:
-        masks: 2-D integer array (0 = background, 1..N = cell labels)
-        num_cells: number of detected cells
+    Dispatch segmentation to the correct venv via subprocess.
+    Returns (masks, num_cells) where masks is a 2-D integer array.
     """
-    img = load_image(image_input)
-
-    diam = diameter if diameter > 0 else None
-
-    normalize_param = {"percentile": [1.0, 99.0], "tile_norm_blocksize": 0}
+    python_exe = get_python_executable(model_type)
 
-    if model_type == "CellposeSAM":
-        if model_name in ("cpsam",):
-            model = models.CellposeModel(gpu=use_gpu, model_type=model_name)
+    # Write input image to a temp file so worker.py can read it
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as img_tmp:
+        img_path = img_tmp.name
+        if isinstance(image_input, np.ndarray):
+            Image.fromarray(image_input.astype(np.uint8)).save(img_path)
         else:
-            model = models.CellposeModel(gpu=use_gpu, pretrained_model=model_name)
-
-        # Pack channels into front of array (same logic as worker.py CellposeSAM path)
-        chan_indices = _parse_channels(channels_text)
-        valid = [c for c in chan_indices if c < img.shape[2]]
-        img_input = np.zeros_like(img)
-        for i, c in enumerate(valid):
-            img_input[:, :, i] = img[:, :, c]
-
-        masks, _, _ = model.eval(
-            img_input,
-            diameter=diam,
-            batch_size=8,
-            resample=True,
-            normalize=normalize_param,
-            flow_threshold=flow_threshold,
-            cellprob_threshold=cellprob_threshold,
-        )[:3]
-
-    else:  # Cellpose3.1
-        if model_name in ("cyto3",):
-            model = models.Cellpose(gpu=use_gpu, model_type=model_name)
-        else:
-            model = models.CellposeModel(gpu=use_gpu, pretrained_model=model_name)
-
-        chan_indices = _parse_channels(channels_text)
-        chan_arg = (chan_indices + [0, 0])[:2]
-
-        masks, _, _ = model.eval(
-            img,
-            diameter=diam,
-            channels=chan_arg,
-            batch_size=8,
-            resample=True,
-            normalize=normalize_param,
-            flow_threshold=flow_threshold,
-            cellprob_threshold=cellprob_threshold,
-        )[:3]
-
-    num_cells = int(masks.max())
-    return masks.astype(np.int32), num_cells
-
-
-def _parse_channels(channels_text: str) -> list[int]:
-    """Parse '0,0' style channel string into a list of ints."""
-    try:
-        parts = [p.strip() for p in str(channels_text).split(",")]
-        result = []
-        for p in parts:
-            val = int(p)
-            result.append(max(0, min(3, val)))
-        return result
-    except Exception:
-        return [0, 0]
+            Image.open(str(image_input)).convert("RGB").save(img_path)
+
+    with tempfile.NamedTemporaryFile(suffix="_mask.png", delete=False) as mask_tmp:
+        mask_path = mask_tmp.name
+
+    command = [
+        python_exe, WORKER,
+        "--image_path", img_path,
+        "--model_type", model_type,
+        "--model_name", model_name,
+        "--diameter", str(diameter),
+        "--flow_threshold", str(flow_threshold),
+        "--cellprob_threshold", str(cellprob_threshold),
+        "--channels", channels_text,
+        "--output_path", mask_path,
+    ]
+
+    result = subprocess.run(
+        command,
+        capture_output=True,
+        text=True,
+    )
+
+    # Print worker stderr so HF Spaces logs show Cellpose progress
+    if result.stderr:
+        for line in result.stderr.strip().splitlines():
+            print(f"[worker] {line}")
+
+    if result.returncode != 0:
+        raise RuntimeError(
+            f"Worker failed (exit {result.returncode}):\n{result.stderr}"
+        )
+
+    output = json.loads(result.stdout.strip())
+    if output.get("status") == "error":
+        raise RuntimeError(output.get("message", "Unknown worker error"))
+
+    num_cells = output.get("num_cells", 0)
+    if num_cells == 0:
+        return np.zeros((1, 1), dtype=np.int32), 0
+
+    # Read the 16-bit PNG mask back as integer array
+    masks = np.array(Image.open(mask_path)).astype(np.int32)
+    return masks, num_cells

setup_venvs.py

@@ -0,0 +1,54 @@
+"""
+Creates venv_v3 (cellpose==3.1.1.2) and venv_v4 (cellpose latest) at Space startup.
+Called once from app.py before the Gradio interface launches.
+"""
+import subprocess
+import sys
+from pathlib import Path
+
+BASE = Path(__file__).parent
+
+VENVS = [
+    ("venv_v3", "requirements_v3.txt"),
+    ("venv_v4", "requirements_v4.txt"),
+]
+
+
+def _python(venv_name: str) -> Path:
+    return BASE / venv_name / "bin" / "python"
+
+
+def setup_venvs():
+    for venv_name, req_file in VENVS:
+        python_path = _python(venv_name)
+        if python_path.exists():
+            print(f"[Setup] {venv_name} already exists — skipping.")
+            continue
+
+        print(f"[Setup] Creating {venv_name} ...")
+        subprocess.run([sys.executable, "-m", "venv", str(BASE / venv_name)], check=True)
+
+        print(f"[Setup] Upgrading pip in {venv_name} ...")
+        subprocess.run([str(python_path), "-m", "pip", "install", "--upgrade", "pip"],
+                       check=True)
+
+        print(f"[Setup] Installing {req_file} into {venv_name} ...")
+        subprocess.run(
+            [str(python_path), "-m", "pip", "install", "-r", str(BASE / req_file)],
+            check=True,
+        )
+        print(f"[Setup] {venv_name} ready.")
+
+
+def get_python_executable(model_type: str) -> str:
+    venv_name = "venv_v3" if model_type == "Cellpose3.1" else "venv_v4"
+    python_path = _python(venv_name)
+    if not python_path.exists():
+        raise RuntimeError(
+            f"venv {venv_name} not found. Call setup_venvs() first."
+        )
+    return str(python_path)
+
+
+if __name__ == "__main__":
+    setup_venvs()

worker.py

@@ -0,0 +1,115 @@
+"""
+Standalone inference script.
+Runs inside either venv_v3 (Cellpose3.1) or venv_v4 (CellposeSAM).
+Writes the integer-label mask as a 16-bit PNG and prints a JSON result to stdout.
+
+Usage:
+    python worker.py --image_path <path> --model_type <CellposeSAM|Cellpose3.1>
+                     --model_name <cpsam|cyto3> --diameter 30
+                     --flow_threshold 0.4 --cellprob_threshold 0.0
+                     --channels 0,0 --output_path <path>
+"""
+import argparse
+import json
+import sys
+
+import numpy as np
+import cv2
+from PIL import Image
+from cellpose import models
+
+
+def load_image(path: str) -> np.ndarray:
+    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+    if img is None:
+        img = np.array(Image.open(path).convert("RGB"))
+        return img.astype(np.uint8)
+    if img.ndim == 2:
+        img = np.stack([img, img, img], axis=-1)
+    elif img.shape[2] == 4:
+        img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
+    else:
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return img.astype(np.uint8)
+
+
+def parse_channels(text: str) -> list:
+    try:
+        return [max(0, min(3, int(p.strip()))) for p in text.split(",")]
+    except Exception:
+        return [0, 0]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--image_path", required=True)
+    parser.add_argument("--model_type", required=True)
+    parser.add_argument("--model_name", required=True)
+    parser.add_argument("--diameter", type=float, default=30.0)
+    parser.add_argument("--flow_threshold", type=float, default=0.4)
+    parser.add_argument("--cellprob_threshold", type=float, default=0.0)
+    parser.add_argument("--channels", default="0,0")
+    parser.add_argument("--output_path", required=True)
+    args = parser.parse_args()
+
+    img = load_image(args.image_path)
+    diam = args.diameter if args.diameter > 0 else None
+    normalize_param = {"percentile": [1.0, 99.0], "tile_norm_blocksize": 0}
+
+    try:
+        if args.model_type == "CellposeSAM":
+            if args.model_name in ("cpsam",):
+                model = models.CellposeModel(gpu=False, model_type=args.model_name)
+            else:
+                model = models.CellposeModel(gpu=False, pretrained_model=args.model_name)
+
+            chan_indices = parse_channels(args.channels)
+            valid = [c for c in chan_indices if c < img.shape[2]]
+            img_input = np.zeros_like(img)
+            for i, c in enumerate(valid):
+                img_input[:, :, i] = img[:, :, c]
+
+            masks, _, _ = model.eval(
+                img_input,
+                diameter=diam,
+                batch_size=8,
+                resample=True,
+                normalize=normalize_param,
+                flow_threshold=args.flow_threshold,
+                cellprob_threshold=args.cellprob_threshold,
+            )[:3]
+
+        else:  # Cellpose3.1
+            if args.model_name in ("cyto3", "nuclei"):
+                model = models.Cellpose(gpu=False, model_type=args.model_name)
+            else:
+                model = models.CellposeModel(gpu=False, pretrained_model=args.model_name)
+
+            chan_arg = (parse_channels(args.channels) + [0, 0])[:2]
+
+            masks, _, _ = model.eval(
+                img,
+                diameter=diam,
+                channels=chan_arg,
+                batch_size=8,
+                resample=True,
+                normalize=normalize_param,
+                flow_threshold=args.flow_threshold,
+                cellprob_threshold=args.cellprob_threshold,
+            )[:3]
+
+        num_cells = int(masks.max())
+
+        # Save integer-label mask as 16-bit PNG
+        mask16 = masks.astype(np.uint16)
+        Image.fromarray(mask16, mode="I;16").save(args.output_path)
+
+        print(json.dumps({"status": "success", "num_cells": num_cells}))
+
+    except Exception as e:
+        print(json.dumps({"status": "error", "message": str(e)}))
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()