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

lavlab_shell-0.5.0/.github/workflows/build.yml

@@ -0,0 +1,31 @@
+name: Build Wheel using Hatch
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  build-wheel:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.12'
+
+      - name: Install build tools
+        run: |
+          python -m pip install --upgrade pip setuptools wheel
+          python -m pip install "virtualenv==20.23.1"
+          python -m pip install --upgrade hatch
+
+      - name: Build wheel with hatch
+        run: hatch build
+      - name: Upload wheel
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheel
+          path: dist/*.whl

{lavlab_shell-0.3.2 → lavlab_shell-0.5.0}/.github/workflows/publish.yml

@@ -28,6 +28,7 @@ jobs:
       - name: Build wheel and sdist inside Docker
         run: |
           docker run --rm \
+            --entrypoint hatch \
             -v "${{ github.workspace }}:/app" \
             -w /app \
             shell-ci:latest build

{lavlab_shell-0.3.2 → lavlab_shell-0.5.0}/.github/workflows/pytest.yml

@@ -7,6 +7,8 @@ on:
 jobs:
   test:
     runs-on: ubuntu-latest
+    env:
+      CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
     steps:
       - uses: actions/checkout@v4
 
@@ -18,11 +20,11 @@ jobs:
       - name: Install test dependencies
         run: |
           python -m pip install --upgrade pip
-          python -m pip install pytest pytest-cov coverage[toml]>=6.2
+          python -m pip install hatch
 
       - name: Run pytest and generate coverage report
         run: |
-          pytest --cov=src --cov-report=xml:coverage.xml
+          hatch run test:cov --cov-report=xml:coverage.xml
 
       - name: Upload coverage artifact
         uses: actions/upload-artifact@v4
@@ -31,14 +33,12 @@ jobs:
           path: ${{ github.workspace }}/coverage.xml
 
       - name: Upload coverage report to Codecov
-        if: ${{ secrets.CODECOV_TOKEN != '' }}
+        if: env.CODECOV_TOKEN != ''
         uses: codecov/codecov-action@v4
         with:
           fail_ci_if_error: true
           files: ${{ github.workspace }}/coverage.xml
-        env:
-          CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
 
       - name: Skip Codecov upload (no token)
-        if: ${{ secrets.CODECOV_TOKEN == '' }}
+        if: env.CODECOV_TOKEN == ''
         run: echo "Skipping Codecov — CODECOV_TOKEN not set."

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: lavlab-shell
-Version: 0.3.2
+Version: 0.5.0
 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
@@ -19,6 +19,8 @@ Requires-Python: <3.13,>=3.10
 Requires-Dist: macenko-pca
 Requires-Dist: monai
 Requires-Dist: numpy
+Requires-Dist: onnxruntime-gpu>=1.17.0; platform_system == 'Linux' and platform_machine != 'arm64'
+Requires-Dist: onnxruntime>=1.17.0; platform_system != 'Linux' or platform_machine == 'arm64'
 Requires-Dist: openslide-python
 Requires-Dist: pyvips
 Requires-Dist: scikit-image

{lavlab_shell-0.3.2 → lavlab_shell-0.5.0}/pyproject.toml

@@ -1,5 +1,5 @@
 [build-system]
-requires = ["hatchling"]
+requires = ["hatch", "hatchling"]
 build-backend = "hatchling.build"
 
 [project]
@@ -34,6 +34,8 @@ dependencies = [
     "openslide-python",
     "macenko-pca",
     "scikit-image",
+    "onnxruntime-gpu>=1.17.0; platform_system == 'Linux' and platform_machine != 'arm64'",
+    "onnxruntime>=1.17.0; platform_system != 'Linux' or platform_machine == 'arm64'",
 ]
 
 [project.optional-dependencies]
@@ -53,8 +55,11 @@ path = "src/shell/__about__.py"
 
 [tool.hatch.build.targets.wheel]
 packages = ["src/shell"]
+exclude = ["src/shell/weights/*.pth"]
+artifacts = ["src/shell/weights/*.onnx"]
 
 [tool.hatch.envs.default]
+installer = "uv"
 dependencies = ["setuptools>=82.0.0"]
 
 [tool.hatch.envs.default.scripts]
@@ -100,6 +105,16 @@ dependencies = [
 build-docs = "mkdocs build"
 serve-docs = "mkdocs serve"
 
+[tool.hatch.envs.export]
+dependencies = [
+    "onnx>=1.16.0",
+    "onnxruntime>=1.17.0",
+    "onnxscript>=0.1.0",
+    "setuptools>=82.0.0",
+]
+[tool.hatch.envs.export.scripts]
+export-model = "python scripts/export_onnx.py"
+
 [[tool.hatch.envs.all.matrix]]
 python = ["3.10", "3.11", "3.12"]
 
@@ -145,8 +160,11 @@ ignore = []
 # Scientific code uses uppercase variable names (Io, W, H) and import aliases
 # (F as torch.nn.functional, BlitzGateway as _BG) that violate pep8-naming.
 "src/shell/preprocessing.py" = ["N803", "N806"]
-"src/shell/inference.py" = ["N812"]
-"src/shell/infer_omero_wsi.py" = ["N814", "F401"]
+"src/shell/inference.py" = ["N812", "E501"]
+"src/shell/infer_omero_wsi.py" = ["N814", "F401", "N806", "E501"]
+"src/shell/infer_wsi.py" = ["N812", "N806", "E501", "T201", "F841", "F821"]
+"src/shell/post_process.py" = ["RUF002", "RUF003", "N806", "B905", "T201", "E501"]
+"src/shell/transforms.py" = ["N803", "N806", "RUF002", "E501"]
 
 [tool.ruff.lint.isort]
 known-first-party = ["shell"]
@@ -157,4 +175,7 @@ docstring-code-format = true
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = ["-ra", "--strict-markers", "--strict-config"]
+filterwarnings = [
+    'ignore:.*torch.*jit.*interface.*:FutureWarning',
+]
 xfail_strict = true

lavlab_shell-0.5.0/scripts/export_onnx.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+"""Export the bundled SegResNetVAE checkpoint to ONNX."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+
+from shell.model import LATEST_MODEL, MODEL_REGISTRY, _resolve_bundled_weights
+
+
+def _build_model_for_export(device: torch.device) -> torch.nn.Module:
+    """Create the reference SegResNetVAE model in eval mode."""
+    from monai.networks.nets import SegResNetVAE
+
+    model = SegResNetVAE(
+        spatial_dims=2,
+        init_filters=16,
+        in_channels=3,
+        out_channels=3,
+        dropout_prob=0.2,
+        norm=("GROUP", {"num_groups": 8}),
+        act=("MISH", {"inplace": True}),
+        input_image_size=(320, 320),
+        vae_nz=256,
+        vae_estimate_std=True,
+    ).to(device)
+    model.eval()
+    return model
+
+
+def export_onnx(
+    checkpoint: str | None = None,
+    *,
+    version: str | None = None,
+    output: str | None = None,
+    device: str = "cpu",
+) -> Path:
+    """Export a bundled checkpoint to ONNX and return the output path."""
+    if checkpoint is None:
+        if version is None:
+            version = LATEST_MODEL
+        checkpoint = str(_resolve_bundled_weights(version))
+
+    dev = torch.device(device)
+    model = _build_model_for_export(dev)
+    state = torch.load(checkpoint, map_location=dev, weights_only=True)
+    model.load_state_dict(state)
+    model.eval()
+
+    if output is None:
+        target = Path(checkpoint).with_suffix(".onnx")
+    else:
+        target = Path(output)
+    target.parent.mkdir(parents=True, exist_ok=True)
+
+    dummy = torch.randn(1, 3, 320, 320, device=dev)
+    torch.onnx.export(
+        model,
+        dummy,
+        str(target),
+        export_params=True,
+        opset_version=17,
+        do_constant_folding=True,
+        input_names=["input"],
+        output_names=["logits"],
+        dynamic_axes={
+            "input": {0: "batch_size"},
+            "logits": {0: "batch_size"},
+        },
+    )
+    return target
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Export the latest SHELL checkpoint to ONNX.")
+    parser.add_argument("--checkpoint", type=str, default=None, help="Path to a .pth checkpoint to export.")
+    parser.add_argument("--version", type=str, default=None, help="Bundled model version to export (e.g. v1).")
+    parser.add_argument("--output", type=str, default=None, help="Destination .onnx path.")
+    parser.add_argument("--device", type=str, default="cpu", help="Device for export (cpu/cuda/mps).")
+    args = parser.parse_args()
+
+    out = export_onnx(args.checkpoint, version=args.version, output=args.output, device=args.device)
+    print(f"Exported ONNX model: {out}")
+
+
+if __name__ == "__main__":
+    main()

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

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

{lavlab_shell-0.3.2 → lavlab_shell-0.5.0}/src/shell/__init__.py

@@ -9,6 +9,14 @@ or simple checks does not eagerly import heavy runtime dependencies
 (like ``torch``, ``monai``, ``pyvips``, or ``omero``).
 """
 
+import warnings
+
 from .__about__ import __version__
 
+warnings.filterwarnings(
+    "ignore",
+    message=r".*torch.*jit.*interface.*",
+    category=FutureWarning,
+)
+
 __all__ = ["__version__"]

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

@@ -105,7 +105,7 @@ def build_parser() -> argparse.ArgumentParser:
     infer_p.add_argument(
         "--target-mpp",
         type=float,
-        default=1.0,
+        default=2.0,
         help="Desired output resolution (um/pixel).",
     )
     infer_p.add_argument(
@@ -125,6 +125,52 @@ def build_parser() -> argparse.ArgumentParser:
         default=None,
         help="Optional: also save the intermediate EHO image.",
     )
+    infer_p.add_argument(
+        "--save-raw",
+        type=str,
+        default=None,
+        help=(
+            "Optional: save raw model predictions as a 3-band uint8 image "
+            "(band 0 = inner / lumen, band 1 = outer / epithelium, "
+            "band 2 = background), scaled to 0/255. "
+            "Saved at the original input resolution. "
+            "Useful for debugging or re-running post-processing offline."
+        ),
+    )
+    infer_p.add_argument(
+        "--profile",
+        type=str,
+        default="best_effort",
+        choices=["best_effort", "precise", "sensitive"],
+        help=(
+            "Post-processing filter profile. "
+            "'best_effort' (default) balances sensitivity and precision; "
+            "'precise' is more conservative; 'sensitive' keeps more predictions."
+        ),
+    )
+    infer_p.add_argument(
+        "--mode",
+        type=str,
+        default="wsi",
+        choices=["wsi", "biopsy", "tile"],
+        help=(
+            "Post-processing mode. "
+            "'wsi' (default): full pipeline with tissue restriction and "
+            "urethra detection. 'biopsy': tissue restriction but no urethra "
+            "detection. 'tile': no tissue mask/urethra; reflect-pads "
+            "predictions before morphological ops. Use 'tile' for individual "
+            "image tiles that lack surrounding context."
+        ),
+    )
+    infer_p.add_argument(
+        "--tile-pad",
+        type=int,
+        default=None,
+        help=(
+            "Reflect-padding in pixels applied on each side in tile mode. "
+            "Defaults to 50%% of the shorter output dimension when not set."
+        ),
+    )
     infer_p.add_argument(
         "--device",
         type=str,
@@ -156,7 +202,7 @@ def build_parser() -> argparse.ArgumentParser:
     omero_p.add_argument(
         "--target-mpp",
         type=float,
-        default=1.0,
+        default=2.0,
         help="Desired output resolution (um/pixel)",
     )
     omero_p.add_argument(
@@ -252,6 +298,17 @@ def build_parser() -> argparse.ArgumentParser:
             "sessions. Default 4."
         ),
     )
+    omero_p.add_argument(
+        "--profile",
+        type=str,
+        default="best_effort",
+        choices=["best_effort", "precise", "sensitive"],
+        help=(
+            "Post-processing filter profile. "
+            "'best_effort' (default) balances sensitivity and precision; "
+            "'precise' is more conservative; 'sensitive' keeps more predictions."
+        ),
+    )
 
     return parser
 
@@ -303,6 +360,10 @@ def main(argv: list[str] | None = None) -> int:
             target_mpp=args.target_mpp,
             mpp=args.mpp,
             save_eho=args.save_eho,
+            save_raw=args.save_raw,
+            profile=args.profile,
+            mode=args.mode,
+            tile_pad=args.tile_pad,
             device=args.device,
         )
 
@@ -351,13 +412,13 @@ def main(argv: list[str] | None = None) -> int:
                 min_tissue_frac=args.min_tissue_frac,
                 prefetch_depth=args.prefetch_depth,
                 num_fetch_workers=args.fetch_workers,
+                profile=args.profile,
             )
             log.debug("DIAG: infer_omero_wsi returned normally")
-        except BaseException as e:
+        except BaseException:
             # 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:")

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

@@ -67,7 +67,7 @@ def _ensure_omero() -> None:
     if _OMERO_LOADED:
         return
     try:
-        import omero  # noqa: F401  (import to ensure runtime available)
+        import omero
         from omero.gateway import BlitzGateway as _BG
         from omero.model import enums as omero_enums
 
@@ -244,7 +244,7 @@ def create_omero_connection(
         # 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)
+            conn._omero_ws_client = client
         except Exception:
             pass
         try:
@@ -1016,7 +1016,7 @@ def _precompute_tile_info(
 
 def _fetch_worker(
     worker_id: int,
-    conn: "BlitzGateway",
+    conn: BlitzGateway,
     work_queue: queue.Queue[Any],
     result_queue: queue.Queue[Any],
     stop_event: threading.Event,
@@ -1175,7 +1175,7 @@ def _parallel_fetch_coordinator(
             conn = create_omero_connection(host, port, username, password)
             conns.append(conn)
             logger.debug("Created OMERO connection for worker %d", i)
-        except Exception as exc:
+        except Exception:
             # Clean up any already-created connections and re-raise.
             for c in conns:
                 try:
@@ -1292,6 +1292,7 @@ def infer_omero_wsi(
     min_tissue_frac: float = 0.01,
     prefetch_depth: int = 8,
     num_fetch_workers: int = 4,
+    profile: str = "best_effort",
 ) -> np.ndarray:
     """Tile-based OMERO inference pipeline.
 
@@ -1633,6 +1634,19 @@ def infer_omero_wsi(
         )
         tiles_tissue = len(tile_schedule)
 
+        # Upscale tissue mask to full resolution for post-processing.
+        th_h, th_w = tissue_mask.shape
+        y_idx = np.clip(
+            (np.arange(out_h) * th_h / out_h).astype(np.int64),
+            0,
+            th_h - 1,
+        )
+        x_idx = np.clip(
+            (np.arange(out_w) * th_w / out_w).astype(np.int64),
+            0,
+            th_w - 1,
+        )
+        tissue_mask_full = tissue_mask[y_idx[:, None], x_idx[None, :]]
         del tissue_mask
         gc.collect()
 
@@ -1691,6 +1705,8 @@ def infer_omero_wsi(
             tile_overlap=tile_overlap,
             roi_size=inference_tile_size,
             sw_overlap=sw_overlap,
+            tissue_mask_full=tissue_mask_full,
+            profile=profile,
             num_fetch_workers=num_fetch_workers,
         )
 
@@ -1739,6 +1755,8 @@ def _run_pipeline(
     tile_overlap: int,
     roi_size: int,
     sw_overlap: float,
+    tissue_mask_full: np.ndarray | None = None,
+    profile: str = "best_effort",
     num_fetch_workers: int = 4,
 ) -> np.ndarray:
     """Producer/consumer pipeline: parallel fetch+EHO → inference in main.
@@ -1820,7 +1838,10 @@ def _run_pipeline(
         except Exception:
             logger.exception("Failed to start OMERO fetch thread; continuing.")
 
-        pred = np.zeros((out_h, out_w), dtype=np.uint8)
+        # Allocate canvases for inner/outer predictions and hematoxylin.
+        inner_pred = np.zeros((out_h, out_w), dtype=bool)
+        outer_pred = np.zeros((out_h, out_w), dtype=bool)
+        hematoxylin_full = np.zeros((out_h, out_w), dtype=np.uint8)
 
         eho_canvas: np.ndarray | None = None
         if save_eho:
@@ -1840,25 +1861,34 @@ def _run_pipeline(
             if eho_canvas is not None:
                 eho_canvas[oy0 : oy0 + oh, ox0 : ox0 + ow] = tile_eho
 
-            tile_pred = run_inference(
+            # Store hematoxylin channel (ch1) for nuclei segmentation.
+            hematoxylin_full[oy0 : oy0 + oh, ox0 : ox0 + ow] = tile_eho[:, :, 1]
+
+            inner_tile, outer_tile = run_inference(
                 tile_eho,
                 model,
                 device,
                 roi_size=(roi_size, roi_size),
                 overlap=sw_overlap,
+                return_raw=True,
             )
             del tile_eho
 
             # Write only the centre-cropped keep region to avoid seams.
-            ph, pw = tile_pred.shape[:2]
+            ph, pw = inner_tile.shape[:2]
             ky1 = min(keep_y0 + keep_h, ph)
             kx1 = min(keep_x0 + keep_w, pw)
-            kept = tile_pred[keep_y0:ky1, keep_x0:kx1]
             out_y0 = oy0 + keep_y0
             out_x0 = ox0 + keep_x0
-            ch, cw = kept.shape[:2]
-            pred[out_y0 : out_y0 + ch, out_x0 : out_x0 + cw] = kept
-            del tile_pred, kept
+            ch = ky1 - keep_y0
+            cw = kx1 - keep_x0
+            inner_pred[out_y0 : out_y0 + ch, out_x0 : out_x0 + cw] = inner_tile[
+                keep_y0:ky1, keep_x0:kx1
+            ]
+            outer_pred[out_y0 : out_y0 + ch, out_x0 : out_x0 + cw] = outer_tile[
+                keep_y0:ky1, keep_x0:kx1
+            ]
+            del inner_tile, outer_tile
 
             processed += 1
             if processed % 25 == 0:
@@ -1876,7 +1906,24 @@ def _run_pipeline(
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-        return _save_results(pred, eho_canvas, save_eho, output_path)
+        # ── Post-processing ──────────────────────────────────────────
+        from shell.post_process import post_process
+
+        logger.info("Running post-processing (profile=%s) …", profile)
+        label_map = post_process(
+            inner_pred,
+            outer_pred,
+            tissue_mask_full
+            if tissue_mask_full is not None
+            else np.ones((out_h, out_w), dtype=bool),
+            hematoxylin_full,
+            profile_name=profile,
+            verbose=True,
+        )
+        del inner_pred, outer_pred, hematoxylin_full, tissue_mask_full
+        gc.collect()
+
+        return _save_results(label_map, eho_canvas, save_eho, output_path)
 
     finally:
         fetch_thread.join(timeout=30)