nnInteractive 2.2.0__tar.gz → 2.3.1__tar.gz

{nninteractive-2.2.0 → nninteractive-2.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nnInteractive
-Version: 2.2.0
+Version: 2.3.1
 Summary: Inference code for nnInteractive
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>
@@ -549,6 +549,8 @@ Link: [![arXiv](https://img.shields.io/badge/arXiv-2503.08373-b31b1b.svg)](https
 # License
 Note that while this repository is available under Apache-2.0 license (see [LICENSE](./LICENSE)), the [model checkpoint](https://huggingface.co/nnInteractive/nnInteractive) is `Creative Commons Attribution Non Commercial Share Alike 4.0`! 
 
+Release model folders ship their own `LICENSE` file whose **first line is the license identifier** (e.g. `CC BY-NC-SA 4.0`); any following lines (such as a link to the full license) are ignored by the tool. At load time this first line is read and exposed as `session.license` so applications can display the model's license prominently. If a checkpoint folder has no `LICENSE` file, the official v1 checkpoint is assumed to be `CC BY-NC-SA 4.0` and any other checkpoint reports `!!MISSING!!`.
+
 # Changelog
 
 ### 1.1.2 - 2025-08-02

{nninteractive-2.2.0 → nninteractive-2.3.1}/nnInteractive/inference/inference_session.py

@@ -79,6 +79,11 @@ class nnInteractiveInferenceSession:
         self.channel_mapping: dict = {}
         self.supports_initial_label: bool = True
         self.supports_zero_shot_label_refinement: bool = True
+        # License of the loaded model checkpoint. Set when the model is loaded
+        # (read from the LICENSE file in the checkpoint folder, or derived for
+        # legacy checkpoints without one). Exposed so GUIs can display it once
+        # the session is initialized. "!!MISSING!!" means the license is unknown.
+        self.license: Optional[str] = None
 
         # image specific
         self.interactions = None  # blosc2.NDArray once initialized
@@ -118,6 +123,31 @@ class nnInteractiveInferenceSession:
             and checkpoint.get("init_args", {}).get("configuration") == "3d_fullres_ps192_bs24"
         )
 
+    @classmethod
+    def _load_license(cls, model_training_output_dir: str, plans: dict, checkpoint: dict) -> str:
+        """Determine the license of the model being loaded.
+
+        Reads the ``LICENSE`` file from the checkpoint folder if present.
+        Expected format: the FIRST line is a short license identifier (e.g.
+        ``CC BY-NC-SA 4.0``); any following lines (URL, full text, …) are for
+        human readers and are ignored. Only the first non-empty line is
+        returned, so ``self.license`` stays a short, displayable string.
+
+        If the folder has no ``LICENSE`` file it is most likely a legacy model:
+        the official v1 checkpoint is CC BY-NC-SA 4.0, anything else is reported
+        as ``"!!MISSING!!"`` so callers (e.g. GUIs) can flag the unknown license.
+        """
+        license_file = join(model_training_output_dir, "LICENSE")
+        if isfile(license_file):
+            with open(license_file, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line:
+                        return line
+        if cls._is_official_checkpoint(plans, checkpoint):
+            return "CC BY-NC-SA 4.0"
+        return "!!MISSING!!"
+
     def _legacy_default_capability(self) -> dict:
         return {
             "supported_interactions": {
@@ -535,7 +565,13 @@ class nnInteractiveInferenceSession:
             dtype=np.float16,
             chunks=(1, *[min(64, s) for s in shape[1:]]),
             blocks=(1, *[min(32, s) for s in shape[1:]]),
-            cparams={"codec": blosc2.Codec.LZ4, "clevel": 5, "nthreads": min(self.torch_n_threads, os.cpu_count())},
+            # Interactions compress better with NOFILTER, which is also faster than SHUFFLE.
+            cparams={
+                "codec": blosc2.Codec.LZ4,
+                "clevel": 5,
+                "filters": [blosc2.Filter.NOFILTER],
+                "nthreads": min(self.torch_n_threads, os.cpu_count()),
+            },
             dparams={"nthreads": 4},
         )
         self._interactions_shape = shape
@@ -604,7 +640,13 @@ class nnInteractiveInferenceSession:
                 dtype=np.float16,
                 chunks=(1, *[min(64, s) for s in self._interactions_shape[1:]]),
                 blocks=(1, *[min(32, s) for s in self._interactions_shape[1:]]),
-                cparams={"codec": blosc2.Codec.LZ4, "clevel": 5, "nthreads": os.cpu_count()},
+                # Interactions compress better with NOFILTER, which is also faster than SHUFFLE.
+                cparams={
+                    "codec": blosc2.Codec.LZ4,
+                    "clevel": 5,
+                    "filters": [blosc2.Filter.NOFILTER],
+                    "nthreads": os.cpu_count(),
+                },
                 dparams={"nthreads": 4},
             )
         self.current_interaction_intensity = 1.0
@@ -1334,6 +1376,16 @@ class nnInteractiveInferenceSession:
         """
         artifacts = self._load_model_artifacts_from_disk(model_training_output_dir, use_fold, checkpoint_name)
         self.initialize_from_loaded_artifacts(artifacts)
+        # With torch.compile the network is compiled lazily on the first forward pass. For a
+        # locally hosted model that lag would otherwise surface on the user's first real
+        # prediction, where it is far more noticeable than during initialization. Trigger the
+        # compilation now with a dummy forward pass so the cost is paid here instead. warmup()
+        # is a no-op when the network is not compiled. The server takes care of its own warmup
+        # explicitly (it shares one compiled network across sessions via
+        # initialize_from_loaded_artifacts), so we only do this on the direct, local entry point.
+        if self.use_torch_compile:
+            print("torch.compile enabled; warming up (compiling) the network now (this is slow once)...")
+            self.warmup()
 
     def _load_model_artifacts_from_disk(
         self,
@@ -1389,12 +1441,11 @@ class nnInteractiveInferenceSession:
         checkpoint = torch.load(
             join(model_training_output_dir, fold_folder, checkpoint_name), map_location=self.device, weights_only=False
         )
-        if self._is_official_checkpoint(plans, checkpoint):
-            print(
-                "License reminder: The official nnInteractive checkpoint is licensed under "
-                "Creative Commons Attribution Non Commercial Share Alike 4.0 (CC BY-NC-SA 4.0). "
-                "See the license note in readme.md (# License)."
-            )
+        self.license = self._load_license(model_training_output_dir, plans, checkpoint)
+        print("=" * 80)
+        print("Model license:")
+        print(self.license)
+        print("=" * 80)
         trainer_name = checkpoint["trainer_name"]
         configuration_name = checkpoint["init_args"]["configuration"]
 
@@ -1440,6 +1491,7 @@ class nnInteractiveInferenceSession:
             "dataset_json": dataset_json,
             "trainer_name": trainer_name,
             "label_manager": plans_manager.get_label_manager(dataset_json),
+            "license": self.license,
         }
 
     def initialize_from_loaded_artifacts(self, artifacts: dict):
@@ -1463,6 +1515,7 @@ class nnInteractiveInferenceSession:
         self.dataset_json = artifacts["dataset_json"]
         self.trainer_name = artifacts["trainer_name"]
         self.label_manager = artifacts["label_manager"]
+        self.license = artifacts["license"]
         if self.use_torch_compile and not isinstance(self.network, OptimizedModule):
             print("Using torch.compile")
             self.network = torch.compile(self.network)

{nninteractive-2.2.0 → nninteractive-2.3.1}/nnInteractive/inference/remote/remote_session.py

@@ -14,6 +14,7 @@ import threading
 import warnings
 from typing import List, Optional, Tuple, Union
 
+import blosc2
 import httpx
 import numpy as np
 import torch
@@ -41,6 +42,16 @@ from nnInteractive.inference.remote._protocol import (
 from nnInteractive.inference.remote.serialization import pack_array, unpack_array
 
 
+def _compression_threads() -> int:
+    """blosc2 thread count for client-side upload compression.
+
+    Full logical CPU count: blosc2 scales measurably onto SMT siblings, so use them all to
+    minimize upload latency. Per-call only (passed to pack_array → compress2), so it never
+    mutates blosc2's global nthreads.
+    """
+    return max(1, os.cpu_count() or 1)
+
+
 class SessionExpiredError(RuntimeError):
     """Raised when the server reports the client's lease no longer exists.
 
@@ -191,6 +202,11 @@ class nnInteractiveRemoteInferenceSession:
         self.preferred_scribble_thickness = caps["preferred_scribble_thickness"]
         self.interaction_decay = caps["interaction_decay"]
         self.INFERENCE_SESSION_VERSION = caps["inference_session_version"]
+        # License of the model loaded on the server. Mirrors
+        # nnInteractiveInferenceSession.license so a GUI can display it
+        # regardless of whether it holds a local or remote session.
+        # "!!MISSING!!" means the server could not determine the license.
+        self.license: Optional[str] = caps.get("license")
 
         self.original_image_shape: Optional[Tuple[int, ...]] = None
         self.target_buffer: Union[np.ndarray, torch.Tensor, None] = None
@@ -284,7 +300,7 @@ class nnInteractiveRemoteInferenceSession:
     def set_image(self, image: np.ndarray, image_properties: Optional[dict] = None) -> None:
         assert image.ndim == 4, f"expected a 4d image as input, got {image.ndim}d. Shape {image.shape}"
         meta = {"image_properties": image_properties or {}}
-        resp = self._post_binary(PATH_SET_IMAGE, meta, pack_array(image))
+        resp = self._post_binary(PATH_SET_IMAGE, meta, pack_array(image, nthreads=_compression_threads()))
         info = resp.json()
         self.original_image_shape = tuple(info["original_image_shape"])
 
@@ -395,7 +411,10 @@ class nnInteractiveRemoteInferenceSession:
             "override_capability_checks": bool(override_capability_checks),
             "interaction_bbox": ([list(b) for b in interaction_bbox] if interaction_bbox is not None else None),
         }
-        resp = self._post_binary(path, meta, pack_array(mask_image))
+        # Interactions (scribble/lasso masks) compress best with NOFILTER; skip auto-selection.
+        resp = self._post_binary(
+            path, meta, pack_array(mask_image, filters=[blosc2.Filter.NOFILTER], nthreads=_compression_threads())
+        )
         self._apply_prediction_response(resp)
 
     def add_initial_seg_interaction(
@@ -420,7 +439,12 @@ class nnInteractiveRemoteInferenceSession:
             "run_prediction": bool(run_prediction),
             "override_capability_checks": bool(override_capability_checks),
         }
-        resp = self._post_binary(PATH_ADD_INITIAL_SEG, meta, pack_array(initial_seg))
+        # Segmentations compress best with NOFILTER; skip auto-selection.
+        resp = self._post_binary(
+            PATH_ADD_INITIAL_SEG,
+            meta,
+            pack_array(initial_seg, filters=[blosc2.Filter.NOFILTER], nthreads=_compression_threads()),
+        )
         self._apply_prediction_response(resp)
 
     # ------------------------------------------------------------------ #

{nninteractive-2.2.0 → nninteractive-2.3.1}/nnInteractive/inference/remote/serialization.py

@@ -47,8 +47,72 @@ _CODEC_ID = {
 _ID_CODEC = {v: k for k, v in _CODEC_ID.items()}
 
 
-def pack_array(arr: np.ndarray, codec: blosc2.Codec = blosc2.Codec.ZSTD, clevel: int = 3) -> bytes:
-    """Serialize a numpy array to a self-describing compressed byte string."""
+# Fraction of each axis used for the center crop that the filter heuristic compresses.
+_SELECT_FILTER_CROP_FRACTION = 0.25
+
+
+def _compress_all(
+    raw: memoryview, total: int, codec: blosc2.Codec, clevel: int, filters: list, nthreads: Optional[int]
+) -> int:
+    """Compressed byte length of ``raw`` under ``filters``, chunked exactly as pack_array does."""
+    extra = {} if nthreads is None else {"nthreads": nthreads}
+    size = 0
+    nchunks = (total + _CHUNK_SIZE - 1) // _CHUNK_SIZE
+    for i in range(nchunks):
+        start = i * _CHUNK_SIZE
+        end = min(start + _CHUNK_SIZE, total)
+        size += len(blosc2.compress2(raw[start:end], codec=codec, clevel=clevel, filters=filters, **extra))
+    return size
+
+
+def _select_filter(arr: np.ndarray, codec: blosc2.Codec, clevel: int, nthreads: Optional[int]) -> "blosc2.Filter":
+    """Pick NOFILTER vs SHUFFLE for ``arr`` by trial-compressing a small centered crop.
+
+    Uses ``compress2`` on the raw bytes — exactly the path pack_array takes — so the decision
+    is consistent with how the whole array is actually compressed. The crop is
+    ``_SELECT_FILTER_CROP_FRACTION`` of each axis (centered), keeping the trial cheap and
+    representative (lands on foreground). Ties go to NOFILTER; any failure falls back to it.
+    """
+    try:
+        crop_shape = [max(1, int(s * _SELECT_FILTER_CROP_FRACTION)) for s in arr.shape]
+        slices = tuple(slice((s - cs) // 2, (s - cs) // 2 + cs) for s, cs in zip(arr.shape, crop_shape))
+        crop = np.ascontiguousarray(arr[slices])
+        raw = memoryview(crop).cast("B")
+        total = raw.nbytes
+
+        best_filter, best_bytes = blosc2.Filter.NOFILTER, None
+        for f in (blosc2.Filter.NOFILTER, blosc2.Filter.SHUFFLE):
+            cb = _compress_all(raw, total, codec, clevel, [f], nthreads)
+            if best_bytes is None or cb < best_bytes:
+                best_bytes, best_filter = cb, f
+        return best_filter
+    except Exception as e:
+        from warnings import warn
+
+        warn(f"_select_filter failed ({e!r}); falling back to NOFILTER.")
+        return blosc2.Filter.NOFILTER
+
+
+def pack_array(
+    arr: np.ndarray,
+    codec: blosc2.Codec = blosc2.Codec.ZSTD,
+    clevel: int = 3,
+    filters: Optional[list] = None,
+    nthreads: Optional[int] = None,
+) -> bytes:
+    """Serialize a numpy array to a self-describing compressed byte string.
+
+    ``filters`` is the blosc2 filter pipeline to apply. If ``None`` (the default), the
+    better of NOFILTER/SHUFFLE is auto-selected by trial-compressing a cheap, representative
+    slab — appropriate for images, whose optimum depends on the data. Callers that already
+    know the optimum (interactions and segmentations compress best with NOFILTER) should pass
+    ``[blosc2.Filter.NOFILTER]`` to skip the selection. The chosen filter is self-describing
+    inside the blosc2 frame, so unpack_array (decompress2) needs no changes.
+
+    ``nthreads`` is the per-call blosc2 thread count for compression. ``None`` (the default)
+    inherits blosc2's global ``nthreads`` (= core count). Passing an explicit value overrides
+    it for this call only, without mutating global state.
+    """
     arr = np.ascontiguousarray(arr)
     dtype_str = arr.dtype.str.lstrip("<>|=").encode("ascii")
     if arr.dtype.byteorder not in ("=", "|", "<"):
@@ -77,11 +141,18 @@ def pack_array(arr: np.ndarray, codec: blosc2.Codec = blosc2.Codec.ZSTD, clevel:
     raw = memoryview(arr).cast("B")
     total = raw.nbytes
     nchunks = (total + _CHUNK_SIZE - 1) // _CHUNK_SIZE
+
+    if filters is None:
+        # Auto-select the better filter from a small centered crop, using the same
+        # compress2 path as below for consistency.
+        filters = [_select_filter(arr, codec, clevel, nthreads)]
+
+    extra = {} if nthreads is None else {"nthreads": nthreads}
     parts = [header, shape_bytes, struct.pack("<I", nchunks)]
     for i in range(nchunks):
         start = i * _CHUNK_SIZE
         end = min(start + _CHUNK_SIZE, total)
-        chunk = blosc2.compress2(raw[start:end], codec=codec, clevel=clevel)
+        chunk = blosc2.compress2(raw[start:end], codec=codec, clevel=clevel, filters=filters, **extra)
         parts.append(struct.pack("<QQ", end - start, len(chunk)))
         parts.append(chunk)
     return b"".join(parts)

{nninteractive-2.2.0 → nninteractive-2.3.1}/nnInteractive/inference/server/app.py

@@ -38,12 +38,15 @@ from __future__ import annotations
 import asyncio
 import json
 import logging
+import numbers
+import os
 import threading
 import time
 import uuid
 from contextlib import asynccontextmanager
 from typing import Optional
 
+import blosc2
 import numpy as np
 import torch
 from fastapi import Depends, FastAPI, HTTPException, Header, Request, Response, status
@@ -73,6 +76,14 @@ from nnInteractive.inference.remote.serialization import pack_array, unpack_arra
 
 logger = logging.getLogger("nninteractive.server")
 
+# Cap a single client's target buffer at 25% of total system RAM. Falls back to 32 GiB
+# of headroom if the system RAM can't be determined.
+try:
+    total_ram = os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE")
+except (ValueError, OSError, AttributeError):
+    total_ram = 32 * 1024**3
+MAX_TARGET_BUFFER_BYTES = int(total_ram * 0.25)
+
 
 class SessionEntry:
     """One client's session plus its bookkeeping."""
@@ -431,11 +442,54 @@ def make_app(
         # Reset so a subsequent call without a prediction can't accidentally re-send a stale region.
         session._last_paste_bbox = None
         return Response(
-            content=pack_array(sub),
+            # Segmentations compress best with NOFILTER; skip auto-selection.
+            content=pack_array(
+                sub, filters=[blosc2.Filter.NOFILTER], nthreads=min(session.torch_n_threads, os.cpu_count())
+            ),
             media_type=CONTENT_TYPE_OCTET_STREAM,
             headers={META_HEADER: json.dumps(meta, separators=(",", ":"))},
         )
 
+    def _parse_target_buffer_request(payload: dict) -> tuple[tuple[int, ...], np.dtype]:
+        if "shape" not in payload:
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="missing required field: shape")
+        if "dtype" not in payload:
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="missing required field: dtype")
+
+        raw_shape = payload["shape"]
+        if not isinstance(raw_shape, list):
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="shape must be a list of positive integers")
+        if len(raw_shape) != 3:
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="shape must be 3D")
+
+        shape = []
+        for dim in raw_shape:
+            if not isinstance(dim, numbers.Integral) or isinstance(dim, bool):
+                raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="shape must contain only integers")
+            if dim <= 0:
+                raise HTTPException(status.HTTP_400_BAD_REQUEST, detail="shape dimensions must be positive")
+            shape.append(dim)
+
+        try:
+            dtype = np.dtype(payload["dtype"])
+        except (TypeError, ValueError) as e:
+            raise HTTPException(status.HTTP_400_BAD_REQUEST, detail=f"invalid dtype: {payload['dtype']!r}") from e
+        # 'b' = bool, 'i' = signed int, 'u' = unsigned int.
+        if dtype.kind not in ("b", "i", "u"):
+            raise HTTPException(
+                status.HTTP_400_BAD_REQUEST,
+                detail=f"unsupported dtype {dtype}: target buffer must be bool or an integer type",
+            )
+
+        nbytes = int(np.prod(shape, dtype=np.uint64)) * dtype.itemsize
+        if nbytes > MAX_TARGET_BUFFER_BYTES:
+            raise HTTPException(
+                status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
+                detail=(f"target buffer would require {nbytes} bytes, " f"limit is {MAX_TARGET_BUFFER_BYTES} bytes"),
+            )
+
+        return tuple(shape), dtype
+
     def _under_session_lock(entry: SessionEntry, fn):
         """Run ``fn(session)`` under the session's lock, converting known errors to HTTP 400.
 
@@ -530,8 +584,7 @@ def make_app(
 
     @app.post(PATH_SET_TARGET_BUFFER, dependencies=[auth])
     def set_target_buffer(payload: dict, entry: SessionEntry = lease) -> dict:
-        shape = tuple(int(x) for x in payload["shape"])
-        dtype = np.dtype(payload["dtype"])
+        shape, dtype = _parse_target_buffer_request(payload)
         buf = np.zeros(shape, dtype=dtype)
 
         def _do(session):
@@ -659,4 +712,5 @@ def _build_capability_snapshot(session: nnInteractiveInferenceSession) -> dict:
         "patch_size": list(cfg.patch_size) if cfg is not None else None,
         "do_autozoom": bool(session.do_autozoom),
         "inference_session_version": session.INFERENCE_SESSION_VERSION,
+        "license": session.license,
     }

{nninteractive-2.2.0 → nninteractive-2.3.1}/nnInteractive.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nnInteractive
-Version: 2.2.0
+Version: 2.3.1
 Summary: Inference code for nnInteractive
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>
@@ -549,6 +549,8 @@ Link: [![arXiv](https://img.shields.io/badge/arXiv-2503.08373-b31b1b.svg)](https
 # License
 Note that while this repository is available under Apache-2.0 license (see [LICENSE](./LICENSE)), the [model checkpoint](https://huggingface.co/nnInteractive/nnInteractive) is `Creative Commons Attribution Non Commercial Share Alike 4.0`! 
 
+Release model folders ship their own `LICENSE` file whose **first line is the license identifier** (e.g. `CC BY-NC-SA 4.0`); any following lines (such as a link to the full license) are ignored by the tool. At load time this first line is read and exposed as `session.license` so applications can display the model's license prominently. If a checkpoint folder has no `LICENSE` file, the official v1 checkpoint is assumed to be `CC BY-NC-SA 4.0` and any other checkpoint reports `!!MISSING!!`.
+
 # Changelog
 
 ### 1.1.2 - 2025-08-02

{nninteractive-2.2.0 → nninteractive-2.3.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nnInteractive"
-version = "2.2.0"
+version = "2.3.1"
 requires-python = ">=3.10"
 description = "Inference code for nnInteractive"
 readme = "readme.md"

{nninteractive-2.2.0 → nninteractive-2.3.1}/readme.md

@@ -311,6 +311,8 @@ Link: [![arXiv](https://img.shields.io/badge/arXiv-2503.08373-b31b1b.svg)](https
 # License
 Note that while this repository is available under Apache-2.0 license (see [LICENSE](./LICENSE)), the [model checkpoint](https://huggingface.co/nnInteractive/nnInteractive) is `Creative Commons Attribution Non Commercial Share Alike 4.0`! 
 
+Release model folders ship their own `LICENSE` file whose **first line is the license identifier** (e.g. `CC BY-NC-SA 4.0`); any following lines (such as a link to the full license) are ignored by the tool. At load time this first line is read and exposed as `session.license` so applications can display the model's license prominently. If a checkpoint folder has no `LICENSE` file, the official v1 checkpoint is assumed to be `CC BY-NC-SA 4.0` and any other checkpoint reports `!!MISSING!!`.
+
 # Changelog
 
 ### 1.1.2 - 2025-08-02