Implement a native function to compute model uncertainty (DAT-1016)

datamint/client_cmd_tools/datamint_inference.py

@@ -57,11 +57,12 @@ def _load_model(project_name: str, model_name: str):
         ) from e
 
 
-def _run_inference(model, file_path: str):
+def _run_inference(model, file_path: str, compute_uncertainty: bool = False):
     from datamint.entities.resource import LocalResource
 
     resource = LocalResource(local_filepath=file_path)
-    predictions = model.predict([resource])
+    params = {'compute_uncertainty': True} if compute_uncertainty else None
+    predictions = model.predict([resource], params=params)
     return resource, (predictions[0] if predictions else [])
 
 
@@ -74,11 +75,14 @@ def _print_predictions(console: Console, predictions: list) -> None:
     table.add_column("Type", style="key")
     table.add_column("Identifier")
     table.add_column("Confidence")
+    table.add_column("Uncertainty")
     for ann in predictions:
         confidence = getattr(ann, 'confiability', None)
         conf_str = f"{confidence:.3f}" if isinstance(confidence, (int, float)) else "-"
+        uncertainty = getattr(ann, 'uncertainty', None)
+        unc_str = f"{uncertainty:.3f}" if isinstance(uncertainty, (int, float)) else "-"
         identifier = getattr(ann, 'identifier', None) or getattr(ann, 'text_value', None) or '-'
-        table.add_row(ann.annotation_type, str(identifier), conf_str)
+        table.add_row(ann.annotation_type, str(identifier), conf_str, unc_str)
     console.print(table)
 
 
@@ -144,7 +148,7 @@ def _execute(args: argparse.Namespace, console: Console) -> int:
         model = _load_model(project_name, args.model_name)
 
     with console.status("[accent]Running inference...[/accent]"):
-        resource, predictions = _run_inference(model, args.file)
+        resource, predictions = _run_inference(model, args.file, compute_uncertainty=args.uncertainty)
 
     _print_predictions(console, predictions)
 
@@ -165,6 +169,8 @@ def _parse_args() -> argparse.Namespace:
                                            # Model registered under a different name than its project
   datamint inference file.png --model-name MyModel --output result.png
                                            # Also save a visualization of the predictions
+  datamint inference file.png --model-name MyModel --uncertainty
+                                           # Also compute a predictive-uncertainty score
 
 More Documentation: https://sonanceai.github.io/datamint-python-api/command_line_tools.html
         """,
@@ -180,6 +186,9 @@ def _parse_args() -> argparse.Namespace:
                         "explicitly to something else).")
     parser.add_argument('--output', type=str, default=None, metavar='PATH',
                         help='Save a rendered visualization of the predictions to this path.')
+    parser.add_argument('--uncertainty', action='store_true', default=False,
+                        help='Also compute a predictive-entropy uncertainty score per prediction '
+                        '(see datamint.utils.uncertainty). Off by default.')
     parser.add_argument('--verbose', action='store_true', default=False, help='Print debug messages.')
 
     return parser.parse_args()

datamint/lightning/trainers/lightning_modules/classification_module.py

@@ -153,9 +153,17 @@ def configure_optimizers(self):
     def predict_default(
         self,
         model_input,
+        compute_uncertainty: bool = False,
         **kwargs: Any,
     ):
-        """Run classification inference, returning :class:`~datamint.entities.annotations.ImageClassification` per resource."""
+        """Run classification inference, returning :class:`~datamint.entities.annotations.ImageClassification` per resource.
+
+        Args:
+            compute_uncertainty: If ``True``, also compute a predictive-entropy
+                uncertainty score (see :mod:`datamint.utils.uncertainty`) and attach
+                it as ``uncertainty`` on each returned annotation.
+        """
+        from datamint.utils.uncertainty import categorical_entropy
 
         transform = self.transform
         if transform is None:
@@ -180,8 +188,12 @@ def predict_default(
                 confidence = float(probs.max(dim=1).values.item())
                 pred_idx = int(logits.argmax(dim=1).item())
                 identifier, value = self.class_names[pred_idx]
+                extra = {}
+                if compute_uncertainty:
+                    extra['uncertainty'] = categorical_entropy(probs)[0].item()
                 all_preds.append([ImageClassification(
                     name=identifier, value=value,
                     confiability=confidence,
+                    **extra,
                 )])
         return all_preds

datamint/lightning/trainers/lightning_modules/segmentation_module.py

@@ -182,11 +182,19 @@ def configure_optimizers(self):
             weight_decay=1e-4,
         )
 
-    def predict_image(self, model_input, **kwargs: Any):
-        """Run segmentation inference, returning :class:`~datamint.entities.annotations.ImageSegmentation` per resource."""
+    def predict_image(self, model_input, compute_uncertainty: bool = False, **kwargs: Any):
+        """Run segmentation inference, returning :class:`~datamint.entities.annotations.ImageSegmentation` per resource.
+
+        Args:
+            compute_uncertainty: If ``True``, also compute a predictive-entropy
+                uncertainty score per class (see :mod:`datamint.utils.uncertainty`,
+                restricted to the pixels predicted as that class) and attach
+                it as ``uncertainty`` on each returned annotation. 
+        """
         import cv2
         import numpy as np
         from datamint.entities.annotations import ImageSegmentation
+        from datamint.utils.uncertainty import segmentation_uncertainty
 
         device = self.inference_device
         self.eval()
@@ -216,9 +224,15 @@ def predict_image(self, model_input, **kwargs: Any):
                         interpolation=cv2.INTER_NEAREST,
                     ) * 255
                     if mask.any():
+                        extra = {}
+                        if compute_uncertainty:
+                            extra['uncertainty'] = segmentation_uncertainty(
+                                probs[0, i], torch.from_numpy(pred[i])
+                            )
                         anns.append(ImageSegmentation(
                             name=name, mask=mask,
                             confiability=class_conf,
+                            **extra,
                         ))
                 all_preds.append(anns)
         return all_preds

datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py

@@ -586,9 +586,14 @@ def _sliding_window_inference(self, volume: Tensor) -> Tensor:
         """
         device = next(self.parameters()).device
         volume = volume.float().to(device)
-        B, C, D, H, W = volume.shape
+        B, C, D0, H0, W0 = volume.shape
         pd, ph, pw = self.patch_crop_size
 
+        pad_d, pad_h, pad_w = max(0, pd - D0), max(0, ph - H0), max(0, pw - W0)
+        if pad_d or pad_h or pad_w:
+            volume = F.pad(volume, (0, pad_w, 0, pad_h, 0, pad_d))
+        _, _, D, H, W = volume.shape
+
         sd = max(1, int(pd * (1 - self.sw_overlap)))
         sh = max(1, int(ph * (1 - self.sw_overlap)))
         sw = max(1, int(pw * (1 - self.sw_overlap)))
@@ -611,7 +616,10 @@ def _tile_starts(total: int, patch: int, step: int) -> list[int]:
                     accum[:, :, d0:d0 + pd, h0:h0 + ph, w0:w0 + pw]   += pred * gauss
                     weights[:, :, d0:d0 + pd, h0:h0 + ph, w0:w0 + pw] += gauss
 
-        return accum / weights.clamp(min=1e-6)
+        result = accum / weights.clamp(min=1e-6)
+        if pad_d or pad_h or pad_w:
+            result = result[:, :, :D0, :H0, :W0]
+        return result
 
     @staticmethod
     def _gaussian_kernel(
@@ -628,3 +636,76 @@ def _g1d(n: int) -> Tensor:
         pd, ph, pw = patch_size
         kernel = _g1d(pd)[:, None, None] * _g1d(ph)[None, :, None] * _g1d(pw)[None, None, :]
         return (kernel / kernel.max()).unsqueeze(0).unsqueeze(0)  # (1, 1, D, H, W)
+
+    # ---- Inference: full-volume prediction -----------------------------------
+
+    def predict_volume(
+        self,
+        model_input,
+        compute_uncertainty: bool = False,
+        uncertainty_top_fraction: float = 0.2,
+        **kwargs: Any,
+    ):
+        """Run sliding-window segmentation inference, returning
+        :class:`~datamint.entities.annotations.VolumeSegmentation` per resource.
+
+        Args:
+            compute_uncertainty: If ``True``, also compute a predictive-entropy
+                uncertainty score per class (see :mod:`datamint.utils.uncertainty`),
+                pooled from per-slice scores to one score per volume, and
+                attach it as ``uncertainty`` on each returned annotation.
+            uncertainty_top_fraction: Fraction of the volume's most-uncertain
+                slices averaged into the per-volume score. Only used when
+                ``compute_uncertainty=True``.
+        """
+        import numpy as np
+        from albumentations.pytorch import ToTensorV2
+        from medimgkit.readers import read_array_normalized
+        from datamint.entities.annotations import VolumeSegmentation
+        from datamint.utils.uncertainty import segmentation_uncertainty, pool_top_k
+
+        transform = self.transform or A.Compose([A.Normalize(), ToTensorV2()])
+        device = self.inference_device
+        self.eval()
+
+        class_names = self.class_names or [f'class_{i}' for i in range(self.num_classes)]
+
+        all_preds: list[list] = []
+        with torch.inference_mode():
+            for res in model_input:
+                res_bytesdata = res.fetch_file_data(auto_convert=False, use_cache=True)
+                img, _ = read_array_normalized(res_bytesdata, return_metainfo=True)  # (N, C, H, W)
+                img = img.transpose(1, 0, 2, 3)  # (C, N, H, W)
+                depth = img.shape[1]
+
+                # Normalize is per-pixel (no spatial context needed across depth), so
+                # applying it slice-by-slice is equivalent to a volume-aware transform.
+                slices = [
+                    transform(image=np.transpose(img[:, d, :, :], (1, 2, 0)))['image']
+                    for d in range(depth)
+                ]
+                volume = torch.stack(slices, dim=1).unsqueeze(0).to(device)  # (1, C, D, H, W)
+
+                logits = self._sliding_window_inference(volume)  # (1, num_classes, D, H, W)
+                probs = torch.sigmoid(logits)
+                pred = logits[0] > 0  # (num_classes, D, H, W)
+
+                anns: list = []
+                for i, name in enumerate(class_names):
+                    mask = pred[i].cpu().numpy().astype(np.int32)  # (D, H, W)
+                    if not mask.any():
+                        continue
+                    extra = {}
+                    if compute_uncertainty:
+                        slice_scores = [
+                            segmentation_uncertainty(probs[0, i, d], pred[i, d])
+                            for d in range(depth)
+                        ]
+                        extra['uncertainty'] = pool_top_k(slice_scores, top_fraction=uncertainty_top_fraction)
+                    anns.append(VolumeSegmentation.from_semantic_segmentation(
+                        segmentation=mask.transpose(1, 2, 0),  # (H, W, D)
+                        class_map=name,
+                        **extra,
+                    ))
+                all_preds.append(anns)
+        return all_preds

datamint/mlflow/flavors/model.py

@@ -188,6 +188,9 @@ def load_context(self, context: PythonModelContext) -> None:
                         f'{context.artifacts=} | {context.model_config=}')
         self._detect_device(context)
         self._move_to_device(self.inference_device)
+
+        #Force a fresh discovery against the code that's actually running now
+        self._invalidate_router()
 
     # ------------------------------------------------------------------
     # Serialization

datamint/mlflow/flavors/prediction_router.py

@@ -26,15 +26,8 @@ class ModeSpec:
 
 
 def bridge_mode(fn: Callable) -> Callable:
-    """Marks a method as a bridge implementation.
-
-    Bridge methods are only auto-registered in the router when their
-    prerequisite mode is already in the registry (see
-    ``PredictionRouter._BRIDGE_MODE_PREREQS``).  Without this marker the
-    router would treat them as ordinary overrides and expose them regardless
-    of whether the prerequisite is implemented.
-    """
-    fn._is_bridge = True  # type: ignore[attr-defined]
+    """Marks a method as a bridge implementation. """
+    fn._is_bridge = True  
     return fn
 
 
@@ -213,11 +206,21 @@ def _resolve_mode(self, model_input: list, params: dict) -> PredictionMode:
             )
         except Exception:
             is_all_image = False
+        try:
+            is_all_volume = all(r.is_volume() for r in model_input)
+        except Exception:
+            is_all_volume = False
 
-        _LOGGER.debug("Parsing prediction mode: '%s' | is_all_image=%s", mode_str, is_all_image)
+        _LOGGER.debug(
+            "Parsing prediction mode: '%s' | is_all_image=%s | is_all_volume=%s",
+            mode_str, is_all_image, is_all_volume,
+        )
 
-        if mode_str == PredictionMode.DEFAULT.value and is_all_image:
-            mode_str = PredictionMode.IMAGE.value
+        if mode_str == PredictionMode.DEFAULT.value:
+            if is_all_image:
+                mode_str = PredictionMode.IMAGE.value
+            elif is_all_volume:
+                mode_str = PredictionMode.VOLUME.value
         try:
             return PredictionMode(mode_str)
         except ValueError:

datamint/utils/uncertainty.py

@@ -0,0 +1,92 @@
+"""Predictive-entropy uncertainty utilities.
+
+Computes how "spread out" a model's own softmax/sigmoid output is. """
+from __future__ import annotations
+
+import torch
+from torch import Tensor
+
+_EPS = 1e-8
+
+
+def categorical_entropy(probs: Tensor, dim: int = -1) -> Tensor:
+    """Normalized Shannon entropy of a categorical probability distribution.
+
+    Args:
+        probs: Probabilities that sum to 1 along ``dim`` (e.g. softmax output).
+        dim: Dimension the distribution lies along.
+
+    Returns:
+        Entropy in ``[0, 1]``, one value per remaining dimension. ``0`` means
+        the distribution is fully concentrated on one class (certain); ``1``
+        means it is uniform across all classes (maximally uncertain).
+    """
+    num_classes = probs.shape[dim]
+    ent = -(probs * torch.log(probs + _EPS)).sum(dim=dim)
+    return ent / torch.log(torch.tensor(float(num_classes), device=probs.device))
+
+
+def binary_entropy(probs: Tensor) -> Tensor:
+    """Normalized binary entropy, element-wise, for sigmoid probabilities.
+
+    Args:
+        probs: Sigmoid probabilities in ``[0, 1]``, any shape.
+
+    Returns:
+        Entropy in ``[0, 1]`` with the same shape as ``probs``. ``0`` at
+        ``p=0`` or ``p=1`` (certain), ``1`` at ``p=0.5`` (maximally uncertain).
+    """
+    ent = -(probs * torch.log(probs + _EPS) + (1 - probs) * torch.log(1 - probs + _EPS))
+    return ent / torch.log(torch.tensor(2.0, device=probs.device))
+
+
+def segmentation_uncertainty(probs: Tensor, pred_mask: Tensor) -> float | None:
+    """Uncertainty score for one class channel of a segmentation prediction.
+
+    Averages per-pixel binary entropy over the pixels the model predicted as
+    foreground for this class. 
+
+    Args:
+        probs: Sigmoid probabilities for this class, any shape (e.g. ``(H,
+            W)`` for a 2-D slice).
+        pred_mask: Boolean/binary mask of the same shape marking pixels the
+            model predicted as this class.
+
+    Returns:
+        Mean entropy over the predicted-foreground pixels, or ``None`` if the
+        model predicted no foreground pixels at all for this class.
+    """
+    fg = pred_mask.bool()
+    if not fg.any():
+        return None
+    ent = binary_entropy(probs)
+    return float(ent[fg].mean())
+
+
+def pool_top_k(scores: list[float | None], top_fraction: float = 0.2) -> float | None:
+    """Collapse many small scores into one representative score.
+
+    Averages only the top ``top_fraction`` highest scores (e.g. the most
+    uncertain slices in a volume).
+
+    Args:
+        scores: Per-item scores; ``None`` entries (e.g. an item with nothing
+            predicted) are skipped.
+        top_fraction: Fraction of the (non-``None``) scores to average, in
+            ``(0, 1]``.
+
+    Returns:
+        The mean of the top ``top_fraction`` scores, or ``None`` if every
+        entry was ``None``.
+    """
+    if not 0 < top_fraction <= 1:
+        raise ValueError(f"top_fraction must be in (0, 1], got {top_fraction}")
+
+    valid = [s for s in scores if s is not None]
+    if not valid:
+        return None
+
+    arr = torch.tensor(valid, dtype=torch.float32)
+    k = max(1, int(len(arr) * top_fraction))
+    top_k = torch.topk(arr, k).values
+    return float(top_k.mean())