From 4a2da016966d19f78a2b8e7f777d5cf869e4dac0 Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Fri, 17 Jul 2026 14:42:26 -0300 Subject: [PATCH] add uncertainty function --- .../client_cmd_tools/datamint_inference.py | 17 +++- .../classification_module.py | 14 ++- .../lightning_modules/segmentation_module.py | 18 +++- .../segmentation_modules/unetrpp.py | 85 ++++++++++++++++- datamint/mlflow/flavors/model.py | 3 + datamint/mlflow/flavors/prediction_router.py | 27 +++--- datamint/utils/uncertainty.py | 92 +++++++++++++++++++ docs/source/command_line_tools.rst | 18 ++++ 8 files changed, 253 insertions(+), 21 deletions(-) create mode 100644 datamint/utils/uncertainty.py diff --git a/datamint/client_cmd_tools/datamint_inference.py b/datamint/client_cmd_tools/datamint_inference.py index b2aea4d5..de3b3439 100644 --- a/datamint/client_cmd_tools/datamint_inference.py +++ b/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() diff --git a/datamint/lightning/trainers/lightning_modules/classification_module.py b/datamint/lightning/trainers/lightning_modules/classification_module.py index 70bac209..44e983a9 100644 --- a/datamint/lightning/trainers/lightning_modules/classification_module.py +++ b/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 diff --git a/datamint/lightning/trainers/lightning_modules/segmentation_module.py b/datamint/lightning/trainers/lightning_modules/segmentation_module.py index 4177e1c3..ef6d4d5c 100644 --- a/datamint/lightning/trainers/lightning_modules/segmentation_module.py +++ b/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 diff --git a/datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py b/datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py index 741294c5..38dd224a 100644 --- a/datamint/lightning/trainers/lightning_modules/segmentation_modules/unetrpp.py +++ b/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 diff --git a/datamint/mlflow/flavors/model.py b/datamint/mlflow/flavors/model.py index 2c40e832..c801c1e7 100644 --- a/datamint/mlflow/flavors/model.py +++ b/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 diff --git a/datamint/mlflow/flavors/prediction_router.py b/datamint/mlflow/flavors/prediction_router.py index 711091b7..2e4be7c4 100644 --- a/datamint/mlflow/flavors/prediction_router.py +++ b/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: diff --git a/datamint/utils/uncertainty.py b/datamint/utils/uncertainty.py new file mode 100644 index 00000000..850e7bec --- /dev/null +++ b/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()) diff --git a/docs/source/command_line_tools.rst b/docs/source/command_line_tools.rst index 815f65b6..a8730524 100644 --- a/docs/source/command_line_tools.rst +++ b/docs/source/command_line_tools.rst @@ -372,4 +372,22 @@ To also save a visualization of the predictions overlaid on the input file, use datamint inference file.png --model-name MyModel --output result.png +Estimating prediction uncertainty ++++++++++++++++++++++++++++++++++ + +Add ``--uncertainty`` to also compute an uncertainty score for each prediction: + +.. code-block:: bash + + datamint inference file.png --model-name MyModel --uncertainty + +The score ranges from 0 (confident) to 1 (maximally uncertain). It is a predictive +entropy computed from the model's own output probabilities in a single forward pass, a +cheap proxy rather than a calibrated estimate. Off by default since it is not needed +for a normal prediction. + +The underlying functions live in :mod:`datamint.utils.uncertainty` and can be called +directly, or reached from the Python API with +``model.predict(resources, params={'compute_uncertainty': True})``. + See all available options by running ``datamint inference --help``. \ No newline at end of file