lavlab-shell 0.3.0__tar.gz → 0.3.2__tar.gz

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: lavlab-shell
-Version: 0.3.0
+Version: 0.3.2
 Summary: SHELL Highlights Epithelium and Lumen Locations — whole-slide H&E segmentation
 Project-URL: Documentation, https://github.com/laviolette-lab/shell#readme
 Project-URL: Issues, https://github.com/laviolette-lab/shell/issues

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/src/shell/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: MIT
 """Version information for shell."""
 
-__version__ = "0.3.0"
+__version__ = "0.3.2"

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/src/shell/cli.py

@@ -12,6 +12,30 @@ Provides two subcommands:
 
 from __future__ import annotations
 
+import os
+
+# Ensure PyTorch initialises its runtime and threads before other native
+# libraries (pyvips, Ice/OMERO) to avoid macOS-level races that can lead
+# to segmentation faults when different C extensions compete for thread
+# runtime ownership.  Set conservative thread counts before importing torch.
+os.environ.setdefault("OMP_NUM_THREADS", "1")
+os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
+os.environ.setdefault("MKL_NUM_THREADS", "1")
+os.environ.setdefault("VECLIB_MAXIMUM_THREADS", "1")
+
+try:
+    # Import torch early so it initialises its native runtime in the main
+    # thread. This is intentional — we only probe __version__ to force
+    # initialisation without performing heavy work here.
+    import torch  # type: ignore
+
+    _ = torch.__version__  # ensure attribute access triggers initialisation
+    TORCH_AVAILABLE = True
+except Exception:
+    # If torch is unavailable we still continue — the program will emit
+    # a clearer error later when trying to build/load the model.
+    TORCH_AVAILABLE = False
+
 import argparse
 import getpass
 import logging
@@ -19,6 +43,9 @@ import sys
 
 from shell.__about__ import __version__
 
+# Module logger — used for diagnostic messages instead of printing directly.
+log = logging.getLogger(__name__)
+
 
 def build_parser() -> argparse.ArgumentParser:
     """Build and return the argument parser."""
@@ -287,35 +314,57 @@ def main(argv: list[str] | None = None) -> int:
         return 0
 
     if args.command == "infer-omero":
+        # Diagnostic prints added to help narrow where a macOS segfault occurs.
+        # These prints are intentionally lightweight and flushed so they appear
+        # even if the process crashes in a C extension later on.
+        log.debug("DIAG: infer-omero selected")
+        log.debug("DIAG: about to import infer_omero_wsi")
         from shell.infer_omero_wsi import infer_omero_wsi
 
+        log.debug("DIAG: imported infer_omero_wsi")
+
         if args.password is None:
             args.password = getpass.getpass("OMERO password: ")
 
-        print(f"Connecting to {args.host}:{args.port} …")
-        infer_omero_wsi(
-            host=args.host,
-            port=args.port,
-            username=args.username,
-            password=args.password,
-            image_id=args.image_id,
-            model_path=args.model_path,
-            model_version=args.model_version,
-            output_path=args.output,
-            target_mpp=args.target_mpp,
-            group_id=args.group_id,
-            save_eho=args.save_eho,
-            no_tissue_crop=args.no_tissue_crop,
-            device=args.device,
-            inference_tile_size=args.inference_tile_size,
-            tile_overlap=args.tile_overlap,
-            sw_overlap=args.sw_overlap,
-            num_bg_tiles=args.num_bg_tiles,
-            min_thumb_size=args.min_thumb_size,
-            min_tissue_frac=args.min_tissue_frac,
-            prefetch_depth=args.prefetch_depth,
-            num_fetch_workers=args.fetch_workers,
-        )
+        print(f"Connecting to {args.host}:{args.port} …", flush=True)
+        log.debug("DIAG: about to call infer_omero_wsi")
+        try:
+            infer_omero_wsi(
+                host=args.host,
+                port=args.port,
+                username=args.username,
+                password=args.password,
+                image_id=args.image_id,
+                model_path=args.model_path,
+                model_version=args.model_version,
+                output_path=args.output,
+                target_mpp=args.target_mpp,
+                group_id=args.group_id,
+                save_eho=args.save_eho,
+                no_tissue_crop=args.no_tissue_crop,
+                device=args.device,
+                inference_tile_size=args.inference_tile_size,
+                tile_overlap=args.tile_overlap,
+                sw_overlap=args.sw_overlap,
+                num_bg_tiles=args.num_bg_tiles,
+                min_thumb_size=args.min_thumb_size,
+                min_tissue_frac=args.min_tissue_frac,
+                prefetch_depth=args.prefetch_depth,
+                num_fetch_workers=args.fetch_workers,
+            )
+            log.debug("DIAG: infer_omero_wsi returned normally")
+        except BaseException as e:
+            # Catch and print Python-level exceptions; note segfaults (native crashes)
+            # will not be caught here, but these diagnostics will show progress up to
+            # the crash point.
+            import sys
+            import traceback
+
+            log.debug("DIAG: infer_omero_wsi raised an exception:")
+            traceback.print_exc()
+            # Re-raise so existing behaviour (non-zero exit) remains.
+            raise
+        log.debug("DIAG: after infer_omero_wsi (should have saved output)")
         print(f"Saved prediction to {args.output}")
         return 0
 

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/src/shell/infer_omero_wsi.py

@@ -48,28 +48,66 @@ logger = logging.getLogger(__name__)
 # ---------------------------------------------------------------------------
 # Optional OMERO dependency gate
 # ---------------------------------------------------------------------------
-try:
-    from omero.gateway import BlitzGateway as _BlitzGateway
-    from omero.model import enums as omero_enums
-
-    PIXEL_TYPES: dict[str, type] = {
-        omero_enums.PixelsTypeint8: np.int8,
-        omero_enums.PixelsTypeuint8: np.uint8,
-        omero_enums.PixelsTypeint16: np.int16,
-        omero_enums.PixelsTypeuint16: np.uint16,
-        omero_enums.PixelsTypeint32: np.int32,
-        omero_enums.PixelsTypeuint32: np.uint32,
-        omero_enums.PixelsTypefloat: np.float32,
-        omero_enums.PixelsTypedouble: np.float64,
-    }
-    _HAS_OMERO = True
-except ImportError:
-    _HAS_OMERO = False
-    PIXEL_TYPES = {}
+# Lazy OMERO loader: defer importing the heavy omero modules until runtime.
+# This avoids importing OMERO (and its C extensions) at module import time
+# which can race with other libraries (e.g. torch / pyvips) and cause
+# platform-specific crashes.
+_OMERO_BG_CLASS = None
+_OMERO_LOADED = False
+PIXEL_TYPES: dict = {}
+
+
+def _ensure_omero() -> None:
+    """Attempt to import OMERO & populate `PIXEL_TYPES` / `_OMERO_BG_CLASS`.
+
+    This is idempotent and safe to call from any runtime location. If the
+    import fails the function leaves `_OMERO_LOADED` as False.
+    """
+    global _OMERO_BG_CLASS, PIXEL_TYPES, _OMERO_LOADED
+    if _OMERO_LOADED:
+        return
+    try:
+        import omero  # noqa: F401  (import to ensure runtime available)
+        from omero.gateway import BlitzGateway as _BG
+        from omero.model import enums as omero_enums
+
+        # Map likely enum attribute names to numpy dtypes, but use getattr
+        # so code is robust to differences across OMERO versions.
+        _mapping = {
+            "PixelsTypeint8": np.int8,
+            "PixelsTypeuint8": np.uint8,
+            "PixelsTypeint16": np.int16,
+            "PixelsTypeuint16": np.uint16,
+            "PixelsTypeint32": np.int32,
+            "PixelsTypeuint32": np.uint32,
+            "PixelsTypefloat": np.float32,
+            "PixelsTypedouble": np.float64,
+        }
+
+        pix = {}
+        for name, dtype in _mapping.items():
+            enum_val = getattr(omero_enums, name, None)
+            if enum_val is not None:
+                pix[enum_val] = dtype
+
+        PIXEL_TYPES.clear()
+        PIXEL_TYPES.update(pix)
+
+        _OMERO_BG_CLASS = _BG
+        _OMERO_LOADED = True
+    except Exception:
+        # Leave PIXEL_TYPES empty and _OMERO_LOADED False on failure.
+        PIXEL_TYPES.clear()
+        _OMERO_BG_CLASS = None
+        _OMERO_LOADED = False
 
 
 def _require_omero() -> None:
-    if not _HAS_OMERO:
+    # Ensure the OMERO runtime is available at call time. This calls the
+    # lazy loader which attempts to import OMERO and populate the runtime
+    # objects. If import fails we raise the usual ImportError.
+    _ensure_omero()
+    if not _OMERO_LOADED:
         msg = (
             "The 'omero' optional dependency is required for OMERO support. "
             "Install it with: pip install shell[omero]"
@@ -142,10 +180,25 @@ def create_omero_connection(
         A connected gateway instance.
     """
     _require_omero()
-    from omero.gateway import BlitzGateway as _BG
+    # Use the lazily-loaded BlitzGateway class populated by _ensure_omero().
+    _BG = _OMERO_BG_CLASS
+    if _BG is None:
+        # Defensive: if the class is not available treat as missing dependency.
+        raise ImportError(
+            "OMERO BlitzGateway class not available; ensure OMERO is installed"
+        )
 
     parsed = urlparse(host)
 
+    # Diagnostic print to help trace macOS crashes during connection setup.
+    try:
+        logger.debug(
+            f"DIAG: create_omero_connection start host={host!r} port={port} scheme={parsed.scheme!r}"
+        )
+    except Exception:
+        # Best-effort diagnostic; do not stop the connection logic if logging fails.
+        pass
+
     if parsed.scheme in ("wss", "ws"):
         import omero  # type: ignore[import-untyped]
 
@@ -166,9 +219,20 @@ def create_omero_connection(
             ws_path,
         )
 
+        try:
+            logger.debug(
+                f"DIAG: creating omero.client with router={router!r} (proto={proto})"
+            )
+        except Exception:
+            pass
+
         client = omero.client(args=["--Ice.Default.Router=" + router])
         try:
             client.createSession(username, password)
+            try:
+                logger.debug("DIAG: omero.client.createSession succeeded")
+            except Exception:
+                pass
         except Exception as exc:
             msg = (
                 f"Failed to create OMERO session via "
@@ -177,18 +241,51 @@ def create_omero_connection(
             raise RuntimeError(msg) from exc
 
         conn = _BG(client_obj=client)
-        if not conn.connect():
+        # If we created a raw omero.client earlier (websocket path) attach it
+        # to the BlitzGateway instance so it can be closed explicitly later.
+        try:
+            setattr(conn, "_omero_ws_client", client)
+        except Exception:
+            pass
+        try:
+            connected = conn.connect()
+            try:
+                logger.debug(f"DIAG: BlitzGateway.connect() returned {connected!r}")
+            except Exception:
+                pass
+        except Exception as exc:
+            # Surface the connect() error with a diagnostic print before raising.
+            try:
+                logger.debug(f"DIAG: BlitzGateway.connect() raised: {exc}")
+            except Exception:
+                pass
+            raise
+
+        if not connected:
             msg = (
                 f"BlitzGateway.connect() failed after session creation "
                 f"({proto}://{ws_host}:{ws_port}{ws_path})"
             )
             raise RuntimeError(msg)
+        try:
+            logger.debug(
+                "DIAG: create_omero_connection returning connected WS BlitzGateway"
+            )
+        except Exception:
+            pass
         return conn
 
     # Standard Ice connection
     clean_host = parsed.hostname or host
     actual_port = parsed.port or port
 
+    try:
+        logger.debug(
+            f"DIAG: creating standard BlitzGateway to {clean_host}:{actual_port} secure={secure}"
+        )
+    except Exception:
+        pass
+
     conn = _BG(
         username,
         password,
@@ -196,12 +293,75 @@ def create_omero_connection(
         port=actual_port,
         secure=secure,
     )
-    if not conn.connect():
+    try:
+        connected = conn.connect()
+        try:
+            logger.debug(f"DIAG: BlitzGateway.connect() returned {connected!r}")
+        except Exception:
+            pass
+    except Exception as exc:
+        try:
+            logger.debug(f"DIAG: BlitzGateway.connect() raised: {exc}")
+        except Exception:
+            pass
+        raise
+
+    if not connected:
         msg = f"Failed to connect to OMERO at {clean_host}:{actual_port}"
         raise RuntimeError(msg)
+
+    try:
+        logger.debug("DIAG: create_omero_connection returning connected BlitzGateway")
+    except Exception:
+        pass
     return conn
 
 
+def _safe_close_blitzgateway(conn: Any) -> None:
+    """Best-effort close of a BlitzGateway and any attached raw omero.client.
+
+    Different OMERO transports / versions expose slightly different cleanup
+    APIs.  Try several likely methods without raising to ensure Ice
+    communicators are destroyed where possible.
+    """
+    try:
+        # Preferred: BlitzGateway.close()
+        if hasattr(conn, "close"):
+            try:
+                conn.close()
+            except Exception:
+                pass
+    except Exception:
+        pass
+
+    try:
+        # Some versions expose disconnect()
+        if hasattr(conn, "disconnect"):
+            try:
+                conn.disconnect()
+            except Exception:
+                pass
+    except Exception:
+        pass
+
+    try:
+        # If we attached a raw omero.client during creation, try to close/destroy it.
+        client_obj = getattr(conn, "_omero_ws_client", None)
+        if client_obj is not None:
+            if hasattr(client_obj, "close"):
+                try:
+                    client_obj.close()
+                except Exception:
+                    pass
+            if hasattr(client_obj, "destroy"):
+                try:
+                    client_obj.destroy()
+                except Exception:
+                    pass
+    except Exception:
+        pass
+
+
 # ---------------------------------------------------------------------------
 # Pixel helpers
 # ---------------------------------------------------------------------------
@@ -622,7 +782,7 @@ def _fetch_thumbnail_rgb(
 
 def _build_tissue_mask(thumb_rgb: np.ndarray) -> np.ndarray:
     """Return a boolean tissue mask (``True`` = tissue) from a thumbnail."""
-    from shell.preprocessing import detect_background
+    from shell.transforms import detect_background
 
     bg_mask = detect_background(thumb_rgb)
     return ~bg_mask
@@ -856,10 +1016,7 @@ def _precompute_tile_info(
 
 def _fetch_worker(
     worker_id: int,
-    host: str,
-    port: int,
-    username: str,
-    password: str,
+    conn: "BlitzGateway",
     work_queue: queue.Queue[Any],
     result_queue: queue.Queue[Any],
     stop_event: threading.Event,
@@ -875,7 +1032,13 @@ def _fetch_worker(
     sy: float,
     stain_params: dict[str, Any],
 ) -> None:
-    """Single fetch worker — owns its own OMERO connection.
+    """Single fetch worker — uses a BlitzGateway connection provided by the
+    main thread.
+
+    The main thread is responsible for creating the OMERO connections and
+    passing a connection object to each worker. Workers must NOT close the
+    shared connection; the coordinator will close connections after workers
+    exit.
 
     Pulls tile descriptors from *work_queue*, fetches pixel data from
     OMERO (using the bulk-fetch fast-path when possible), applies EHO
@@ -884,12 +1047,12 @@ def _fetch_worker(
     The bounded *result_queue* provides back-pressure: if inference falls
     behind, workers block on ``put`` until space is available.
     """
-    from shell.preprocessing import apply_eho_chunked
+    from shell.transforms import apply_eho_chunked
 
-    conn = None
+    # Connection is provided by the caller (created in the main/coordinator
+    # thread). Worker owns a RawPixelsStore instance derived from it.
     store = None
     try:
-        conn = create_omero_connection(host, port, username, password)
         store = conn.c.sf.createRawPixelsStore()
         store.setPixelsId(pixels_id, False)
         store.setResolutionLevel(best_level)
@@ -952,16 +1115,14 @@ def _fetch_worker(
         logger.exception("Fetch worker %d encountered an error", worker_id)
         result_queue.put(exc)
     finally:
+        # Close only the RawPixelsStore opened by this worker.
         if store is not None:
             try:
                 store.close()
             except Exception:
                 pass
-        if conn is not None:
-            try:
-                conn.close()
-            except Exception:
-                pass
+        # Do NOT close the BlitzGateway connection here; the coordinator
+        # (creator) is responsible for closing it.
 
 
 def _parallel_fetch_coordinator(
@@ -986,10 +1147,10 @@ def _parallel_fetch_coordinator(
 ) -> None:
     """Start *num_workers* fetch threads, each with its own OMERO session.
 
-    All tile descriptors are placed on a shared work queue.  Workers pull
-    tiles, fetch pixels, apply EHO, and place results on *result_queue*.
-    When the work queue is exhausted a ``_SENTINEL`` is placed on
-    *result_queue* so the consumer knows all tiles have been produced.
+    The coordinator now creates the BlitzGateway connections in the main
+    coordinator thread and passes a dedicated connection to each worker.
+    This avoids importing / initialising the OMERO native runtime inside
+    worker threads (which can trigger macOS-level races and segfaults).
     """
     work_queue: queue.Queue[Any] = queue.Queue()
     stop_event = threading.Event()
@@ -1006,16 +1167,31 @@ def _parallel_fetch_coordinator(
         n_tiles,
     )
 
+    # Pre-create one BlitzGateway connection per worker in this (main)
+    # coordinator thread to avoid initialising OMERO / Ice inside workers.
+    conns: list[Any] = []
+    for i in range(num_workers):
+        try:
+            conn = create_omero_connection(host, port, username, password)
+            conns.append(conn)
+            logger.debug("Created OMERO connection for worker %d", i)
+        except Exception as exc:
+            # Clean up any already-created connections and re-raise.
+            for c in conns:
+                try:
+                    c.close()
+                except Exception:
+                    pass
+            logger.exception("Failed to create OMERO connection for worker %d", i)
+            raise
+
     workers: list[threading.Thread] = []
     for i in range(num_workers):
         t = threading.Thread(
             target=_fetch_worker,
             kwargs={
                 "worker_id": i,
-                "host": host,
-                "port": port,
-                "username": username,
-                "password": password,
+                "conn": conns[i],
                 "work_queue": work_queue,
                 "result_queue": result_queue,
                 "stop_event": stop_event,
@@ -1043,6 +1219,14 @@ def _parallel_fetch_coordinator(
     result_queue.put(_SENTINEL)
     logger.info("All fetch workers finished.")
 
+    # Close all pre-created BlitzGateway connections used by the workers.
+    # Use the best-effort helper to avoid leaving Ice communicators open.
+    for c in conns:
+        try:
+            _safe_close_blitzgateway(c)
+        except Exception:
+            pass
+
 
 # ---------------------------------------------------------------------------
 # Save helpers
@@ -1186,8 +1370,19 @@ def infer_omero_wsi(
     np.ndarray
         ``(H, W)`` uint8 label map.
     """
-    _require_omero()
+    try:
+        logger.debug(
+            f"DIAG: infer_omero_wsi start host={host!r} port={port} image_id={image_id} "
+            f"inference_tile_size={inference_tile_size} num_fetch_workers={num_fetch_workers}"
+        )
+    except Exception:
+        pass
 
+    _require_omero()
+    try:
+        logger.debug("DIAG: _require_omero() returned")
+    except Exception:
+        pass
     # Sanity-check: overlap must leave room for a meaningful core.
     if tile_overlap < 0:
         tile_overlap = 0
@@ -1200,6 +1395,36 @@ def infer_omero_wsi(
         )
         tile_overlap = inference_tile_size // 4
 
+    # Ensure PyTorch and the model are initialised in the main thread BEFORE
+    # creating OMERO connections / starting any worker threads. This avoids
+    # macOS-level races between native runtimes (torch / libvips / Ice).
+    try:
+        import torch  # type: ignore
+
+        from shell.model import build_model
+    except Exception:
+        # If torch or model scaffolding is not available we'll let the
+        # existing _require_omero() / create_omero_connection surface an
+        # appropriate error later. Continue to attempt OMERO connection.
+        torch = None  # type: ignore
+        build_model = None  # type: ignore
+
+    if build_model is not None:
+        if device == "auto":
+            try:
+                if torch is not None and torch.cuda.is_available():
+                    device = "cuda"
+                elif torch is not None and torch.backends.mps.is_available():
+                    device = "mps"
+                else:
+                    device = "cpu"
+            except Exception:
+                device = "cpu"
+        logger.info("Loading model (device=%s) …", device)
+        model = build_model(model_path, device, model_version=model_version)
+    else:
+        model = None  # type: ignore
+
     conn = create_omero_connection(host, port, username, password)
 
     try:
@@ -1316,7 +1541,7 @@ def infer_omero_wsi(
         # ------------------------------------------------------------------
         # 3. Fetch a few background tiles at full resolution for Io
         # ------------------------------------------------------------------
-        from shell.preprocessing import compute_background_intensity
+        from shell.transforms import compute_background_intensity
 
         bg_tile_coords = _pick_background_tile_coords(
             tissue_mask,
@@ -1459,19 +1684,19 @@ def infer_omero_wsi(
             out_h=out_h,
             out_w=out_w,
             prefetch_depth=prefetch_depth,
-            model_path=model_path,
-            model_version=model_version,
+            model=model,
             device=device,
             save_eho=save_eho,
             output_path=output_path,
             tile_overlap=tile_overlap,
+            roi_size=inference_tile_size,
             sw_overlap=sw_overlap,
             num_fetch_workers=num_fetch_workers,
         )
 
     except Exception:
         try:
-            conn.close()
+            _safe_close_blitzgateway(conn)
         except Exception:
             pass
         raise
@@ -1507,12 +1732,12 @@ def _run_pipeline(
     out_h: int,
     out_w: int,
     prefetch_depth: int,
-    model_path: str | None,
-    model_version: str | None,
+    model: Any,
     device: str,
     save_eho: str | None,
     output_path: str,
     tile_overlap: int,
+    roi_size: int,
     sw_overlap: float,
     num_fetch_workers: int = 4,
 ) -> np.ndarray:
@@ -1566,26 +1791,34 @@ def _run_pipeline(
         name="omero-fetch-coordinator",
         daemon=True,
     )
-    fetch_thread.start()
-
     try:
-        # torch is imported ONLY here, in the main thread.  The fetch
-        # thread never touches torch.
-        import torch
+        # torch is imported here for device checks used by run_inference.
+        # The model itself is provided by the caller (loaded earlier in the
+        # main infer_omero_wsi function) and must be passed via the `model`
+        # parameter. Starting the fetch thread now is safe because the model
+        # has already been initialised on the main thread.
+        import torch  # type: ignore
 
         from shell.inference import run_inference
-        from shell.model import build_model
 
         if device == "auto":
-            if torch.cuda.is_available():
-                device = "cuda"
-            elif torch.backends.mps.is_available():
-                device = "mps"
-            else:
+            try:
+                if torch.cuda.is_available():
+                    device = "cuda"
+                elif torch.backends.mps.is_available():
+                    device = "mps"
+                else:
+                    device = "cpu"
+            except Exception:
                 device = "cpu"
 
-        logger.info("Loading model (device=%s) …", device)
-        model = build_model(model_path, device, model_version=model_version)
+        # Start the fetch thread immediately; the heavy model runtime was
+        # already initialised by the caller before calling _run_pipeline.
+        try:
+            fetch_thread.start()
+            logger.info("Started OMERO fetch thread.")
+        except Exception:
+            logger.exception("Failed to start OMERO fetch thread; continuing.")
 
         pred = np.zeros((out_h, out_w), dtype=np.uint8)
 
@@ -1611,6 +1844,7 @@ def _run_pipeline(
                 tile_eho,
                 model,
                 device,
+                roi_size=(roi_size, roi_size),
                 overlap=sw_overlap,
             )
             del tile_eho
@@ -1649,6 +1883,6 @@ def _run_pipeline(
         if fetch_thread.is_alive():
             logger.warning("Fetch thread did not exit within 30 s; continuing cleanup.")
         try:
-            conn.close()
+            _safe_close_blitzgateway(conn)
         except Exception:
             pass

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/src/shell/infer_wsi.py

@@ -24,6 +24,7 @@ import logging
 import os
 import warnings
 from types import ModuleType
+from typing import cast
 
 import numpy as np
 
@@ -147,8 +148,18 @@ def _load_image(
         raise ValueError(msg) from exc
 
 
-def _vips_to_rgb_numpy(vips_img: pyvips.Image) -> np.ndarray:
-    """Convert a pyvips image to (H, W, 3) uint8 RGB numpy array."""
+def _vips_to_rgb_numpy(vips_img: pyvips.Image | np.ndarray) -> np.ndarray:
+    """Convert a pyvips image or numpy array to (H, W, 3) uint8 RGB numpy array."""
+    # If the caller accidentally passes a numpy array (e.g. from an openslide
+    # fallback), accept it and normalise to (H, W, 3) uint8.
+    if isinstance(vips_img, np.ndarray):
+        arr = vips_img
+        if arr.ndim == 2:
+            arr = np.stack([arr, arr, arr], axis=-1)
+        elif arr.ndim == 3 and arr.shape[2] == 4:
+            arr = arr[..., :3]
+        return arr.astype(np.uint8)
+
     bands = vips_img.bands
     if bands == 1:
         vips_img = vips_img.bandjoin([vips_img, vips_img])
@@ -251,11 +262,11 @@ def preprocess_wsi(
     needs_scaling = not (abs(scale_x - 1.0) < 1e-6 and abs(scale_y - 1.0) < 1e-6)
 
     if source == "vips":
-        vips_img = img_or_vips
+        vips_img = cast(pyvips.Image, img_or_vips)
         if needs_scaling:
-            vips_img = vips_img.resize(
-                1.0 / scale_x, vscale=1.0 / scale_y, kernel="lanczos3"
-            )
+            # Use positional vscale/kernel to satisfy the pyvips stubs and
+            # avoid type-checker complaints about keyword-only overloads.
+            vips_img = vips_img.resize(1.0 / scale_x, 1.0 / scale_y, "lanczos3")
         image_np = _vips_to_rgb_numpy(vips_img)
         del vips_img
     else:
@@ -263,8 +274,10 @@ def preprocess_wsi(
         image_np = img_or_vips
         if needs_scaling:
             vips_tmp = pyvips.Image.new_from_array(image_np)
-            vips_tmp = vips_tmp.resize(
-                1.0 / scale_x, vscale=1.0 / scale_y, kernel="lanczos3"
+            # cast to Image for the type-checker and use positional args for
+            # the same reason as above.
+            vips_tmp = cast(pyvips.Image, vips_tmp).resize(
+                1.0 / scale_x, 1.0 / scale_y, "lanczos3"
             )
             image_np = _vips_to_rgb_numpy(vips_tmp)
             del vips_tmp

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/src/shell/inference.py

@@ -21,6 +21,7 @@ in the output, saving downstream processing.
 from __future__ import annotations
 
 import gc
+import os
 import warnings
 
 import numpy as np
@@ -39,7 +40,7 @@ from torch.amp import autocast
 # ---------------------------------------------------------------------------
 # Default inference parameters
 # ---------------------------------------------------------------------------
-VAL_ROI_SIZE: tuple[int, int] = (512, 512)
+VAL_ROI_SIZE: tuple[int, int] = (2048, 2048)
 VAL_SW_BATCH: int = 16
 VAL_SW_OVERLAP: float = 0.25
 
@@ -63,9 +64,7 @@ def _build_normalize_pipeline() -> Compose:
                 b_max=1.0,
                 clip=True,
             ),
-            NormalizeIntensityd(
-                keys="image", nonzero=True, channel_wise=True
-            ),
+            NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
             ScaleIntensityd(keys=["image"]),
         ]
     )
@@ -91,7 +90,7 @@ def run_inference(
     device: torch.device | str = "cpu",
     *,
     roi_size: tuple[int, int] = VAL_ROI_SIZE,
-    sw_batch_size: int = VAL_SW_BATCH,
+    sw_batch_size: int | None = None,
     overlap: float = VAL_SW_OVERLAP,
     tissue_mask: np.ndarray | None = None,
 ) -> np.ndarray:
@@ -101,7 +100,7 @@ def run_inference(
     :param model: trained SegResNetVAE in eval mode.
     :param device: computation device.
     :param roi_size: sliding-window patch size.
-    :param sw_batch_size: number of patches per forward pass.
+    :param sw_batch_size: number of patches per forward pass (None => auto).
     :param overlap: fraction of overlap between sliding-window patches.
         Values > 0 enable Gaussian importance weighting so that
         overlapping patch centres contribute more than edges, which
@@ -130,6 +129,53 @@ def run_inference(
     if pad_h or pad_w:
         img_t = F.pad(img_t, padding, "reflect")
 
+    # Heuristic: choose a sliding-window batch size automatically when not provided.
+    def _choose_sw_batch_size(roi: tuple[int, int], device_obj: torch.device) -> int:
+        """Heuristic selection based on ROI area and device capabilities.
+
+        Baseline: VAL_SW_BATCH for 512x512 on a mid-range GPU. Scale batch size
+        inversely with ROI area and up-weight for devices with more memory
+        (e.g. macOS unified memory / MPS).
+        """
+        base_area = 512 * 512
+        roi_area = max(1, roi[0] * roi[1])
+        area_ratio = base_area / roi_area
+
+        # Try to estimate system RAM (in GB) on POSIX systems; conservative fallback.
+        total_ram_gb: float | None = None
+        try:
+            pages = os.sysconf("SC_PHYS_PAGES")
+            page_size = os.sysconf("SC_PAGE_SIZE")
+            total_ram_gb = (pages * page_size) / (1024**3)
+        except Exception:
+            total_ram_gb = None
+
+        # Device factor: raise batch size for GPUs, raise more for large unified RAM (MPS).
+        if device_obj.type == "cuda":
+            device_factor = 2.0
+        elif device_obj.type == "mps":
+            # macOS unified memory benefits from larger batches when system RAM is large.
+            if total_ram_gb is not None:
+                device_factor = max(1.5, min(8.0, total_ram_gb / 8.0))
+            else:
+                device_factor = 3.0
+        else:
+            # CPU: be conservative and scale with available RAM if known.
+            if total_ram_gb is not None:
+                device_factor = max(0.5, min(2.0, total_ram_gb / 32.0))
+            else:
+                device_factor = 0.5
+
+        batch = max(1, int(VAL_SW_BATCH * area_ratio * device_factor))
+        # Clamp to avoid enormous batches on pathological inputs.
+        return min(max(batch, 1), 1024)
+
+    local_sw_batch = (
+        sw_batch_size
+        if sw_batch_size is not None
+        else _choose_sw_batch_size(roi_size, device_obj)
+    )
+
     amp_device = "cuda" if device_obj.type == "cuda" else "cpu"
     with (
         torch.inference_mode(),
@@ -145,17 +191,69 @@ def run_inference(
         logits = sliding_window_inference(
             img_t,
             roi_size,
-            sw_batch_size,
+            local_sw_batch,
             model,
             overlap=overlap,
             sw_device=device_obj,
             device=torch.device("cpu"),
             mode=blend_mode,
         )
+        # MONAI's sliding_window_inference may return a Tensor, a tuple/list
+        # (e.g. when the model returns multiple outputs), or a dict. Normalise
+        # to a single Tensor here so downstream code and the type-checker see
+        # a consistent type.
+        if not isinstance(logits, torch.Tensor):
+            # Tuple / list -> first element is expected to be the logits Tensor.
+            if isinstance(logits, (tuple, list)):
+                if len(logits) > 0 and isinstance(logits[0], torch.Tensor):
+                    logits = logits[0]
+                else:
+                    # Fallback: try to coerce the first element to a tensor.
+                    logits = torch.as_tensor(logits[0])
+            elif isinstance(logits, dict):
+                # Take the first tensor-like value from the dict.
+                found = False
+                for v in logits.values():
+                    if isinstance(v, torch.Tensor):
+                        logits = v
+                        found = True
+                        break
+                if not found:
+                    # Coerce the first value to a tensor as a last resort.
+                    first_val = next(iter(logits.values()))
+                    logits = torch.as_tensor(first_val)
+            else:
+                # If it's some other type, attempt a coercion to Tensor.
+                logits = torch.as_tensor(logits)
     del img_t
 
     # Remove padding
     if pad_h or pad_w:
+        # Ensure we have a Tensor before accessing shape (some MONAI
+        # wrappers may return a tuple/dict). Normalise common container
+        # types to a single Tensor.
+        if not isinstance(logits, torch.Tensor):
+            if (
+                isinstance(logits, (tuple, list))
+                and len(logits) > 0
+                and isinstance(logits[0], torch.Tensor)
+            ):
+                logits = logits[0]
+            elif isinstance(logits, dict):
+                # take first tensor-like value
+                found = False
+                for v in logits.values():
+                    if isinstance(v, torch.Tensor):
+                        logits = v
+                        found = True
+                        break
+                if not found:
+                    # fall back to coercing first value
+                    first_val = next(iter(logits.values()))
+                    logits = torch.as_tensor(first_val)
+            else:
+                logits = torch.as_tensor(logits)
+
         _, _, ph, pw = logits.shape
         logits = logits[
             :, :, padding[2] : ph - padding[3], padding[0] : pw - padding[1]
@@ -164,12 +262,15 @@ def run_inference(
     # ------------------------------------------------------------------
     # Post-processing: sigmoid + threshold + background comparison
     # ------------------------------------------------------------------
+    # Ensure logits is a Tensor before applying sigmoid.
+    if not isinstance(logits, torch.Tensor):
+        logits = torch.as_tensor(logits)
     probs = torch.sigmoid(logits)
     del logits
 
     inner_prob = probs[0, 0]  # channel 0 = epithelium / inner
     outer_prob = probs[0, 1]  # channel 1 = stroma / outer
-    bg_prob = probs[0, 2]     # channel 2 = background
+    bg_prob = probs[0, 2]  # channel 2 = background
     del probs
 
     inner_mask = (inner_prob > 0.5) & (inner_prob > bg_prob)

{lavlab_shell-0.3.0 → lavlab_shell-0.3.2}/src/shell/transforms.py

@@ -22,15 +22,40 @@ from __future__ import annotations
 
 import logging
 
+import macenko_pca
 import numpy as np
-import pyvips as pv
 from monai.data import MetaTensor
 from monai.transforms import MapTransform, Resize
 from scipy.ndimage import label, uniform_filter
 from scipy.ndimage import sum as ndimage_sum
-from skimage.color import rgb2gray
 
-import macenko_pca
+
+# Lightweight local replacement for skimage.color.rgb2gray to avoid an
+# optional dependency and to keep typing simple. This implements the
+# standard luminance weights and expects input in the [0, 1] float range
+# (matching skimage behavior). It returns a 2D array of floats in [0, 1].
+def rgb2gray(image: np.ndarray) -> np.ndarray:
+    if image.ndim == 2:
+        return image
+    # Expect HWC or CHW: prefer HWC (common in this module)
+    if image.ndim == 3 and image.shape[-1] == 3:
+        # Use common luminance weights
+        weights = np.array([0.2989, 0.5870, 0.1140], dtype=image.dtype)
+        return np.dot(image[..., :3], weights)
+    # Fallback: reduce any trailing channel dimension
+    if image.ndim == 3:
+        return np.mean(image, axis=-1)
+    # Unexpected shape: coerce to 2D via mean
+    return np.mean(image, axis=-1)
+
+
+def _get_pyvips():
+    try:
+        import pyvips as pv
+    except Exception as e:
+        raise RuntimeError("pyvips is required for this operation") from e
+    return pv
+
 
 # ---------------------------------------------------------------------------
 # Resolution-scaling transforms
@@ -40,14 +65,16 @@ import macenko_pca
 class Rescaled(MapTransform):
     """Rescale images from *in_res* to *out_res* microns-per-pixel."""
 
-    def __init__(self, keys, in_res=0.46, out_res=1.37):
+    def __init__(self, keys, in_res=0.46, out_res=2):
         super().__init__(keys)
         self.in_res = in_res
         self.out_res = out_res
 
-    def __call__(self, data, in_res=None, resize_kwargs={}):
+    def __call__(self, data, in_res=None, resize_kwargs=None):
         if in_res is None:
             in_res = self.in_res
+        if resize_kwargs is None:
+            resize_kwargs = {}
         for key in self.keys:
             image = data[key]
             height, width = image.shape[1:]
@@ -55,9 +82,7 @@ class Rescaled(MapTransform):
             new_width = width * in_res // self.out_res
 
             interpolation = (
-                "linear"
-                if len(image.shape) == 3 and image.shape[0] == 3
-                else "nearest"
+                "linear" if len(image.shape) == 3 and image.shape[0] == 3 else "nearest"
             )
 
             if "mode" not in resize_kwargs:
@@ -99,9 +124,8 @@ class LoadImageAtScaled(MapTransform):
         for key in self.key_iterator(d):
             filepath = d[key]
             try:
-                image_pv = pv.Image.new_from_file(
-                    str(filepath), **self.reader_kwargs
-                )
+                pv = _get_pyvips()
+                image_pv = pv.Image.new_from_file(str(filepath), **self.reader_kwargs)
             except Exception as e:
                 raise RuntimeError(
                     f"Failed to load image {filepath} with pyvips: {e}"
@@ -118,9 +142,7 @@ class LoadImageAtScaled(MapTransform):
             ):
                 scale_factor = self.in_res / self.target_res
                 if abs(scale_factor - 1.0) > 1e-6:
-                    image_pv = image_pv.resize(
-                        scale_factor, kernel=self.resize_kernel
-                    )
+                    image_pv = image_pv.resize(scale_factor, kernel=self.resize_kernel)
                     applied_scale_factor = scale_factor
 
             img_array = image_pv.numpy()
@@ -141,9 +163,7 @@ class LoadImageAtScaled(MapTransform):
             if self.target_res is not None:
                 meta_info["target_resolution"] = self.target_res
             if applied_scale_factor != 1.0:
-                meta_info[
-                    "applied_geometric_scale_factor"
-                ] = applied_scale_factor
+                meta_info["applied_geometric_scale_factor"] = applied_scale_factor
 
             d[key] = MetaTensor(img_array, meta=meta_info)
         return d
@@ -172,9 +192,7 @@ class TissueMaskd(MapTransform):
             bg_mask = detect_background(image)
             tissue_mask = ~bg_mask
             meta_info = {"original_channel_dim": "no_channel"}
-            data[f"{key}_tissue_mask"] = MetaTensor(
-                tissue_mask, meta=meta_info
-            )
+            data[f"{key}_tissue_mask"] = MetaTensor(tissue_mask, meta=meta_info)
         return data
 
 
@@ -209,14 +227,21 @@ class MarcenkoDeconvolutiond(MapTransform):
                 tissue_mask = resizer(tissue_mask)
                 tissue_mask = tissue_mask[0]
 
+        # If a tissue_mask was provided as a tensor/MetaTensor, coerce to numpy
+        if tissue_mask is not None:
+            tissue_mask = np.asarray(tissue_mask)
+
         if tissue_mask is None:
             bg_mask = detect_background(image)
             tissue_mask = np.ascontiguousarray(~bg_mask).astype(bool)
 
+        # Ensure tissue_mask is a NumPy boolean array for downstream processing
+        tissue_mask = np.asarray(tissue_mask).astype(bool)
+
         return macenko_pca.rgb_separate_stains_macenko_pca(
             np.ascontiguousarray(image),
             None,
-            mask_out=~tissue_mask.astype(bool),
+            mask_out=~tissue_mask,
         )
 
     def deconv(self, image, tissue_mask=None):
@@ -225,16 +250,12 @@ class MarcenkoDeconvolutiond(MapTransform):
 
         w_est = self.get_w_est(image, tissue_mask=tissue_mask)
 
-        image_her = macenko_pca.color_deconvolution(
-            np.ascontiguousarray(image), w_est
-        )
+        image_her = macenko_pca.color_deconvolution(np.ascontiguousarray(image), w_est)
 
         hematox_index = macenko_pca.find_stain_index(
             self.stain_color_map["hematoxylin"], w_est
         )
-        eosin_index = macenko_pca.find_stain_index(
-            self.stain_color_map["eosin"], w_est
-        )
+        eosin_index = macenko_pca.find_stain_index(self.stain_color_map["eosin"], w_est)
 
         image_her = image_her[..., [hematox_index, eosin_index, 2]]
 
@@ -341,9 +362,7 @@ def detect_background(
     del diff, nz, cmax
 
     gray_q = (
-        0.299 * image_np[..., 0]
-        + 0.587 * image_np[..., 1]
-        + 0.114 * image_np[..., 2]
+        0.299 * image_np[..., 0] + 0.587 * image_np[..., 1] + 0.114 * image_np[..., 2]
     ).astype(np.uint8) >> 4
     entropy_map = np.zeros(gray_q.shape, dtype=np.float32)
     for b in range(16):
@@ -409,25 +428,17 @@ def compute_background_intensity(
 
     if np.sum(bg_for_io) >= 100:
         io_val = float(
-            np.clip(
-                np.mean(image_np[bg_for_io], dtype=np.float64), 200, 255
-            )
+            np.clip(np.mean(image_np[bg_for_io], dtype=np.float64), 200, 255)
         )
     elif np.sum(bg_mask) >= 100:
-        io_val = float(
-            np.clip(
-                np.mean(image_np[bg_mask], dtype=np.float64), 200, 255
-            )
-        )
+        io_val = float(np.clip(np.mean(image_np[bg_mask], dtype=np.float64), 200, 255))
     else:
         io_val = 240.0
     del bg_for_io
     return io_val
 
 
-def estimate_stain_params(
-    rgb: np.ndarray, bg_mask: np.ndarray | None = None
-) -> dict:
+def estimate_stain_params(rgb: np.ndarray, bg_mask: np.ndarray | None = None) -> dict:
     """Estimate Macenko colour-deconvolution parameters from an RGB image."""
     if bg_mask is None:
         bg_mask = detect_background(rgb)
@@ -612,25 +623,16 @@ class EHOd(MapTransform):
             est_image = image
             est_bg = bg_mask
             if image.size > 16777216:  # > 4096^2 pixels
-                logging.info(
-                    "Down-sampling image for stain-parameter estimation"
-                )
+                logging.info("Down-sampling image for stain-parameter estimation")
                 scale = 4
                 h, w = image.shape[:2]
                 small_h, small_w = h // scale, w // scale
-                resizer = Resize(
-                    spatial_size=(small_h, small_w), mode="linear"
-                )
+                resizer = Resize(spatial_size=(small_h, small_w), mode="linear")
                 est_image = (
-                    resizer(image.transpose((2, 0, 1)))
-                    .numpy()
-                    .transpose((1, 2, 0))
+                    resizer(image.transpose((2, 0, 1))).numpy().transpose((1, 2, 0))
                 )
                 if est_bg is not None:
-                    est_bg = (
-                        resizer(est_bg[None].astype(np.float32)).numpy()[0]
-                        > 0.5
-                    )
+                    est_bg = resizer(est_bg[None].astype(np.float32)).numpy()[0] > 0.5
 
             params = estimate_stain_params(
                 np.ascontiguousarray(est_image).astype(np.uint8),

lavlab_shell-0.3.0/src/shell/preprocessing.py

@@ -1,21 +0,0 @@
-# SPDX-FileCopyrightText: 2024-present barrettMCW <mjbarrett@mcw.edu>
-#
-# SPDX-License-Identifier: MIT
-"""
-Colour preprocessing for H&E whole-slide images.
-
-This module re-exports the lower-level helper functions from
-:mod:`shell.transforms` so that existing call-sites (e.g. the OMERO
-tile-by-tile pipeline in :mod:`shell.infer_omero_wsi`) continue to work
-without modification.
-
-For the full set of MONAI dictionary transforms see
-:mod:`shell.transforms`.
-"""
-
-from shell.transforms import (  # noqa: F401
-    apply_eho_chunked,
-    compute_background_intensity,
-    detect_background,
-    estimate_stain_params,
-)