From 72f733c1b4739a6cd867480b7cc942fb5ac33f54 Mon Sep 17 00:00:00 2001
From: PierWSL2 <piero.policastro@gmail.com>
Date: Mon, 27 Jul 2026 08:44:35 +0200
Subject: [PATCH 1/6] feature add: MC-Dropout uncertainty estimation

---
 .gitignore                                    |  10 +
 tests/test_mc_dropout_units.py                | 172 ++++++++++++++++++
 .../models/architecture/bayesian_wrapper.py   | 125 +++++++++++++
 tiatoolbox/models/engine/mc_dropout.py        |  80 ++++++++
 tiatoolbox/models/engine/uncertainty.py       |  36 ++++
 5 files changed, 423 insertions(+)
 create mode 100644 tests/test_mc_dropout_units.py
 create mode 100644 tiatoolbox/models/architecture/bayesian_wrapper.py
 create mode 100644 tiatoolbox/models/engine/mc_dropout.py
 create mode 100644 tiatoolbox/models/engine/uncertainty.py

diff --git a/.gitignore b/.gitignore
index 66c072da50..244a724da7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -119,3 +119,13 @@ ENV/
 
 # output zarr generated
 *.zarr
+
+.opencode/
+graphify-out/
+AGENTS.md
+outputs/
+
+# Local scratch / interview artifacts (not part of the feature)
+/tiny.svs
+/main.py
+/uv.lock
diff --git a/tests/test_mc_dropout_units.py b/tests/test_mc_dropout_units.py
new file mode 100644
index 0000000000..a321b8647f
--- /dev/null
+++ b/tests/test_mc_dropout_units.py
@@ -0,0 +1,172 @@
+"""Unit tests for the three MC-Dropout building blocks, no network/WSI needed.
+
+- tiatoolbox.models.engine.mc_dropout      -> mc_dropout_mode, has_dropout_layers,
+                                               inject_dropout_before_classifier
+- tiatoolbox.models.engine.uncertainty     -> decompose_uncertainty
+- tiatoolbox.models.architecture.bayesian_wrapper -> BayesianModelWrapper
+
+Run all:
+    pytest tests/test_mc_dropout_units.py -v -s
+
+Run one file's tests only, e.g. just mc_dropout:
+    pytest tests/test_mc_dropout_units.py -k mc_dropout -v -s
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from torch import nn
+
+from tiatoolbox.models.architecture.bayesian_wrapper import BayesianModelWrapper
+from tiatoolbox.models.engine.mc_dropout import (
+    has_dropout_layers,
+    inject_dropout_before_classifier,
+    mc_dropout_mode,
+)
+from tiatoolbox.models.engine.uncertainty import decompose_uncertainty
+
+# ----------------------------------------------------------------------------
+# Small dummy models for testing the MC Dropout utilities without needing a full CNN or WSI.
+# ----------------------------------------------------------------------------
+
+
+class DummyCNNModel(nn.Module):
+    """Minimal stand-in for tiatoolbox's CNNModel: feat_extract -> classifier."""
+
+    def __init__(self, with_dropout: bool = False) -> None:
+        super().__init__()
+        self.feat_extract = nn.Conv2d(3, 4, kernel_size=3)
+        self.bn = nn.BatchNorm2d(4)
+        self.pool = nn.AdaptiveAvgPool2d((1, 1))
+        linear = nn.Linear(4, 2)
+        self.classifier = nn.Sequential(nn.Dropout(0.5), linear) if with_dropout else linear
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.bn(self.feat_extract(x))
+        x = self.pool(x).flatten(1)
+        return self.classifier(x)
+
+    @staticmethod
+    def infer_batch(model: "DummyCNNModel", batch_data: torch.Tensor, device: str = "cpu") -> torch.Tensor:
+        """Mimic tiatoolbox's CNNModel.infer_batch(model, images, device) signature."""
+        with torch.no_grad():
+            logits = model(batch_data.to(device))
+            return torch.softmax(logits, dim=-1)
+
+
+# ----------------------------------------------------------------------------
+# mc_dropout.py
+# ----------------------------------------------------------------------------
+def test_has_dropout_layers_true_and_false() -> None:
+    assert has_dropout_layers(DummyCNNModel(with_dropout=True)) is True
+    assert has_dropout_layers(DummyCNNModel(with_dropout=False)) is False
+
+
+def test_inject_dropout_before_classifier_wraps_linear() -> None:
+    model = DummyCNNModel(with_dropout=False)
+    model = inject_dropout_before_classifier(model, p=0.3)
+
+    assert isinstance(model.classifier, nn.Sequential)
+    assert isinstance(model.classifier[0], nn.Dropout)
+    assert model.classifier[0].p == 0.3
+    assert has_dropout_layers(model) is True
+
+
+def test_inject_dropout_before_classifier_is_idempotent() -> None:
+    model = DummyCNNModel(with_dropout=False)
+    model = inject_dropout_before_classifier(model, p=0.3)
+    n_children_before = len(list(model.classifier.children()))
+
+    model = inject_dropout_before_classifier(model, p=0.9)  # should be a no-op
+
+    assert len(list(model.classifier.children())) == n_children_before
+    assert model.classifier[0].p == 0.3  # untouched, not overwritten to 0.9
+
+
+def test_mc_dropout_mode_activates_only_dropout() -> None:
+    model = DummyCNNModel(with_dropout=True)
+    model.eval()  # everything off to start
+
+    with mc_dropout_mode(model) as m:
+        assert m.classifier[0].training is True  # Dropout -> train mode
+        assert m.bn.training is False  # BatchNorm must stay in eval mode
+
+    # After the context exits, everything is restored to eval
+    assert model.classifier[0].training is False
+    assert model.bn.training is False
+
+
+def test_mc_dropout_mode_restores_previous_training_state() -> None:
+    model = DummyCNNModel(with_dropout=True)
+    model.train()  # whole model was in train mode before entering
+
+    with mc_dropout_mode(model):
+        pass
+
+    assert model.training is True  # restored, not forced to eval afterwards
+    assert model.bn.training is True
+
+
+# ----------------------------------------------------------------------------
+# uncertainty.py
+# ----------------------------------------------------------------------------
+def test_decompose_uncertainty_shapes_and_keys() -> None:
+    torch.manual_seed(0)
+    t, batch, n_classes = 20, 5, 3
+    logits = torch.randn(t, batch, n_classes)
+    mc_probs = torch.softmax(logits, dim=-1)
+
+    stats = decompose_uncertainty(mc_probs, class_dim=-1)
+
+    assert set(stats) == {"mean_probs", "epistemic", "aleatoric", "total"}
+    assert stats["mean_probs"].shape == (batch, n_classes)
+    assert stats["epistemic"].shape == (batch,)
+    assert stats["aleatoric"].shape == (batch,)
+    assert stats["total"].shape == (batch,)
+    assert torch.allclose(stats["total"], stats["epistemic"] + stats["aleatoric"])
+
+
+def test_decompose_uncertainty_zero_variance_gives_zero_epistemic() -> None:
+    """If every MC sample is identical, epistemic uncertainty must be exactly 0."""
+    single_pass = torch.tensor([[0.7, 0.2, 0.1]])
+    mc_probs = single_pass.unsqueeze(0).repeat(10, 1, 1)  # [T=10, B=1, C=3], all identical
+
+    stats = decompose_uncertainty(mc_probs, class_dim=-1)
+
+    assert torch.allclose(stats["epistemic"], torch.zeros_like(stats["epistemic"]))
+    assert stats["aleatoric"].item() > 0  # data itself is still uncertain
+
+
+# ----------------------------------------------------------------------------
+# bayesian_wrapper.py
+# ----------------------------------------------------------------------------
+def test_bayesian_wrapper_raises_without_dropout() -> None:
+    model = DummyCNNModel(with_dropout=False)
+    try:
+        BayesianModelWrapper(base_model=model, n_samples=5, inject_dropout=False)
+    except ValueError as exc:
+        assert "Dropout" in str(exc)
+    else:
+        msg = "Expected ValueError when wrapping a model with no Dropout layers."
+        raise AssertionError(msg)
+
+
+def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
+    model = DummyCNNModel(with_dropout=True)
+    model.eval()
+
+    wrapper = BayesianModelWrapper(base_model=model, n_samples=8, class_dim=-1)
+    dummy_batch = torch.rand(4, 3, 8, 8)
+
+    mean_probs = wrapper.infer_batch(wrapper, dummy_batch, device="cpu")
+
+    assert isinstance(mean_probs, np.ndarray)
+    assert mean_probs.shape == (4, 2)  # [B, C]
+    assert len(wrapper.uncertainty_stats) == 1
+
+    batch_stats = wrapper.uncertainty_stats[0]
+    for key in ("mean_probs", "aleatoric", "epistemic", "total", "mc_probs"):
+        assert key in batch_stats
+    assert batch_stats["mc_probs"].shape == (8, 4, 2)  # [T, B, C]
+
diff --git a/tiatoolbox/models/architecture/bayesian_wrapper.py b/tiatoolbox/models/architecture/bayesian_wrapper.py
new file mode 100644
index 0000000000..c84a29c5bc
--- /dev/null
+++ b/tiatoolbox/models/architecture/bayesian_wrapper.py
@@ -0,0 +1,125 @@
+"""Transparent Bayesian wrapper for TIAToolbox models compatible with infer_batch.
+
+The engine (`EngineABC.infer_patches`) calls
+`infer_batch(self.model, batch_data["image"], device=self.device)`, hence the
+static-method signature `(model, images, device)` is preserved so the engine
+can keep passing the wrapper as the first argument. Per-batch uncertainty
+statistics are accumulated on `wrapper.uncertainty_stats` (a list, in
+inference order) so that downstream code — after `predictor.run()` finishes —
+can concatenate them along axis 0 to obtain one entry per patch, aligned
+with the `coordinates` / `probabilities` arrays written by the engine.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+from tiatoolbox.models.engine.mc_dropout import (
+    has_dropout_layers,
+    inject_dropout_before_classifier,
+    mc_dropout_mode,
+)
+from tiatoolbox.models.engine.uncertainty import decompose_uncertainty
+
+
+class BayesianModelWrapper(torch.nn.Module):
+    """Wrap a TIAToolbox model and add Monte Carlo Dropout at `infer_batch`.
+
+    The wrapped model must contain at least one Dropout layer. For pretrained
+    models that ship without Dropout (e.g. `resnet18-kather100k`,
+    `alexnet-kather100k`), pass `inject_dropout=True` (default) and a
+    `nn.Dropout` is inserted before the linear classifier automatically,
+    via `inject_dropout_before_classifier`.
+
+    After each `infer_batch` call, the per-batch uncertainty decomposition
+    (mean_probs, aleatoric, epistemic, total, mc_probs) is appended to
+    `self.uncertainty_stats` in inference order, so it can be aligned with
+    the `coordinates` array written to the engine's output `.zarr`.
+    """
+
+    def __init__(
+        self,
+        base_model: torch.nn.Module,
+        n_samples: int = 30,
+        class_dim: int = -1,
+        inject_dropout: bool = True,
+        dropout_p: float = 0.2,
+    ) -> None:
+        super().__init__()
+        if not has_dropout_layers(base_model):
+            if not inject_dropout:
+                msg = (
+                    "Model has no Dropout layers: MC Dropout is not applicable. "
+                    "Pass inject_dropout=True or add an nn.Dropout before wrapping."
+                )
+                raise ValueError(msg)
+            base_model = inject_dropout_before_classifier(base_model, p=dropout_p)
+            if not has_dropout_layers(base_model):
+                msg = (
+                    "inject_dropout_before_classifier failed to insert "
+                    "any Dropout layer into the provided model."
+                )
+                raise RuntimeError(msg)
+
+        self.base_model = base_model
+        self.n_samples = n_samples
+        self.class_dim = class_dim
+        # Per-batch stats populated by `infer_batch`, in inference order.
+        # Concatenate along axis 0 after `run()` to align with the engine's
+        # `coordinates` / `probabilities` outputs.
+        self.uncertainty_stats: list[dict[str, np.ndarray]] = []
+
+    def __getattr__(self, name: str) -> Any:
+        # Delegate undefined attributes to the wrapped model (e.g. preproc_func,
+        # postproc_func, class_dict) so the engine finds them via _get_model_attr.
+        # `__getattr__` is invoked only when normal lookup fails, so by the time
+        # we reach here `self.base_model` is already accessible normally.
+        try:
+            return super().__getattr__(name)
+        except AttributeError:
+            pass
+        base_model = self.__dict__.get("_modules", {}).get("base_model")
+        if base_model is not None:
+            return getattr(base_model, name)
+        msg = f"{type(self).__name__!s} has no attribute {name!r}"
+        raise AttributeError(msg)
+
+    @staticmethod
+    def infer_batch(
+        model: "BayesianModelWrapper",
+        batch_data: torch.Tensor,
+        device: str = "cpu",
+    ) -> np.ndarray:
+        """Run `n_samples` MC forward passes under `mc_dropout_mode`.
+
+        Returns the MC-mean per-class probabilities of shape [B, C] (numpy),
+        and appends a per-batch uncertainty dict to `model.uncertainty_stats`.
+        """
+        wrapper = model  # the engine passes the wrapper as the first argument
+        base_infer = wrapper.base_model.infer_batch
+
+        samples: list[torch.Tensor] = []
+        with mc_dropout_mode(wrapper.base_model):
+            for _ in range(wrapper.n_samples):
+                sample = base_infer(wrapper.base_model, batch_data, device)
+                samples.append(torch.as_tensor(sample))
+
+        mc_probs = torch.stack(samples, dim=0)  # [T, B, C]
+        stats = decompose_uncertainty(mc_probs, class_dim=wrapper.class_dim)
+
+        wrapper.uncertainty_stats.append(
+            {
+                "mean_probs": stats["mean_probs"].cpu().numpy(),
+                "aleatoric": stats["aleatoric"].cpu().numpy(),
+                "epistemic": stats["epistemic"].cpu().numpy(),
+                "total": stats["total"].cpu().numpy(),
+                "mc_probs": mc_probs.cpu().numpy(),
+            }
+        )
+
+        return stats["mean_probs"].to(
+            next(wrapper.base_model.parameters()).device
+        ).cpu().numpy()
diff --git a/tiatoolbox/models/engine/mc_dropout.py b/tiatoolbox/models/engine/mc_dropout.py
new file mode 100644
index 0000000000..8c2c35a52f
--- /dev/null
+++ b/tiatoolbox/models/engine/mc_dropout.py
@@ -0,0 +1,80 @@
+"""Utilities to enable Monte Carlo Dropout without touching BatchNorm."""
+
+from __future__ import annotations
+
+from torch import nn
+
+_DROPOUT_TYPES = (
+    nn.Dropout,
+    nn.Dropout1d,
+    nn.Dropout2d,
+    nn.Dropout3d,
+    nn.AlphaDropout,
+    nn.FeatureAlphaDropout,
+)
+
+
+class mc_dropout_mode:
+    """Context manager: enable only Dropout layers, everything else stays in eval."""
+
+    def __init__(self, model: nn.Module) -> None:
+        self.model = model
+        # Save training state of ALL modules (not just Dropout) so BatchNorm
+        # and the root module are restored correctly on exit, since
+        # `self.model.eval()` puts the whole tree in eval mode.
+        self._prev_states: dict[nn.Module, bool] = {}
+
+    def __enter__(self) -> nn.Module:
+        self._prev_states = {m: m.training for m in self.model.modules()}
+
+        self.model.eval()
+        for module in self.model.modules():
+            if isinstance(module, _DROPOUT_TYPES):
+                module.train()
+
+        return self.model
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        for module, was_training in self._prev_states.items():
+            module.train(was_training)
+
+        return False  # do not suppress exceptions
+
+
+def has_dropout_layers(model: nn.Module) -> bool:
+    """Check if the model has any dropout layers."""
+    return any(isinstance(m, _DROPOUT_TYPES) for m in model.modules())
+
+
+def inject_dropout_before_classifier(
+    model: nn.Module,
+    p: float = 0.2,
+) -> nn.Module:
+    """Insert a `nn.Dropout` before the model's linear classifier head.
+
+    Pretrained TIAToolbox patch classifiers (e.g. `resnet18-kather100k`,
+    `alexnet-kather100k`) ship with no Dropout layers, which would make
+    Monte Carlo Dropout a no-op. This injects a single `nn.Dropout` between
+    the global-average-pool features and the linear classifier, the
+    canonical placement for MC Dropout on a ResNet/AlexNet classifier.
+
+    Idempotent: if a `nn.Dropout` already immediately precedes the
+    classifier it is left untouched.
+    """
+    head = getattr(model, "classifier", None)
+    if head is None:
+        msg = (
+            "Expected a CNNModel-like module with a `classifier` head, "
+            f"got {type(model).__name__}."
+        )
+        raise TypeError(msg)
+    if isinstance(head, nn.Sequential) and len(head) >= 2 and isinstance(head[0], nn.Dropout):
+        return model
+    if not isinstance(head, nn.Linear):
+        msg = (
+            "Expected `model.classifier` to be an nn.Linear, "
+            f"got {type(head).__name__}."
+        )
+        raise TypeError(msg)
+    model.classifier = nn.Sequential(nn.Dropout(p=p), head)
+    return model
diff --git a/tiatoolbox/models/engine/uncertainty.py b/tiatoolbox/models/engine/uncertainty.py
new file mode 100644
index 0000000000..900635cd92
--- /dev/null
+++ b/tiatoolbox/models/engine/uncertainty.py
@@ -0,0 +1,36 @@
+"""
+Uncertainty Quantification with Monte Carlo (MC Dropout).
+Gal & Ghahramani (2015), used by Policastro (2020) WSI analysis project.
+"""
+
+from __future__ import annotations
+
+import torch
+
+
+def decompose_uncertainty(
+    mc_probs: torch.Tensor, class_dim: int = -1
+) -> dict[str, torch.Tensor]:
+    """Decompose MC-Dropout uncertainty into epistemic and aleatoric terms.
+
+    Args:
+        mc_probs: tensor [T, ..., C] with T Monte Carlo forward passes,
+            each already softmax-normalized along class_dim.
+        class_dim: class axis in mc_probs (excluding T).
+
+    Returns:
+        dict: mean_probs, epistemic, aleatoric, total uncertainty.
+    """
+    mean_probs = mc_probs.mean(dim=0)          # [..., C]
+    var_probs = mc_probs.var(dim=0)            # [..., C]
+
+    epistemic = var_probs.mean(dim=class_dim)  # variance across T passes
+    aleatoric = (mean_probs * (1 - mean_probs)).mean(dim=class_dim)
+    total = epistemic + aleatoric
+
+    return {
+        "mean_probs": mean_probs,
+        "epistemic": epistemic,
+        "aleatoric": aleatoric,
+        "total": total,
+    }
\ No newline at end of file

From d644b7b8eaa21bd2a1d56c177f96d97b96d50d7e Mon Sep 17 00:00:00 2001
From: PierWSL2 <piero.policastro@gmail.com>
Date: Mon, 27 Jul 2026 09:54:00 +0200
Subject: [PATCH 2/6] fix: apply ruff formatting and precommit tests

---
 tests/test_mc_dropout_units.py                | 52 ++++++++++----
 .../models/architecture/bayesian_wrapper.py   | 68 ++++++++++++++++---
 tiatoolbox/models/engine/mc_dropout.py        | 68 +++++++++++++++++--
 tiatoolbox/models/engine/uncertainty.py       | 16 +++--
 4 files changed, 168 insertions(+), 36 deletions(-)

diff --git a/tests/test_mc_dropout_units.py b/tests/test_mc_dropout_units.py
index a321b8647f..a0fe372c5f 100644
--- a/tests/test_mc_dropout_units.py
+++ b/tests/test_mc_dropout_units.py
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import numpy as np
+import pytest
 import torch
 from torch import nn
 
@@ -27,28 +28,51 @@
 from tiatoolbox.models.engine.uncertainty import decompose_uncertainty
 
 # ----------------------------------------------------------------------------
-# Small dummy models for testing the MC Dropout utilities without needing a full CNN or WSI.
+# Small dummy models for testing the MC Dropout utilities without needing a
+# full CNN or WSI.
 # ----------------------------------------------------------------------------
 
 
 class DummyCNNModel(nn.Module):
     """Minimal stand-in for tiatoolbox's CNNModel: feat_extract -> classifier."""
 
-    def __init__(self, with_dropout: bool = False) -> None:
+    def __init__(self, *, with_dropout: bool = False) -> None:
+        """Initialize :class:`DummyCNNModel`.
+
+        Args:
+            with_dropout (bool):
+                If True, prepend a ``nn.Dropout(0.5)`` to the classifier head.
+
+        """
         super().__init__()
         self.feat_extract = nn.Conv2d(3, 4, kernel_size=3)
         self.bn = nn.BatchNorm2d(4)
         self.pool = nn.AdaptiveAvgPool2d((1, 1))
         linear = nn.Linear(4, 2)
-        self.classifier = nn.Sequential(nn.Dropout(0.5), linear) if with_dropout else linear
+        self.classifier = (
+            nn.Sequential(nn.Dropout(0.5), linear) if with_dropout else linear
+        )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Pass input data through the model.
+
+        Args:
+            x (torch.Tensor):
+                Model input.
+
+        Returns:
+            torch.Tensor:
+                The output logits after passing through the model.
+
+        """
         x = self.bn(self.feat_extract(x))
         x = self.pool(x).flatten(1)
         return self.classifier(x)
 
     @staticmethod
-    def infer_batch(model: "DummyCNNModel", batch_data: torch.Tensor, device: str = "cpu") -> torch.Tensor:
+    def infer_batch(
+        model: DummyCNNModel, batch_data: torch.Tensor, device: str = "cpu"
+    ) -> torch.Tensor:
         """Mimic tiatoolbox's CNNModel.infer_batch(model, images, device) signature."""
         with torch.no_grad():
             logits = model(batch_data.to(device))
@@ -59,11 +83,13 @@ def infer_batch(model: "DummyCNNModel", batch_data: torch.Tensor, device: str =
 # mc_dropout.py
 # ----------------------------------------------------------------------------
 def test_has_dropout_layers_true_and_false() -> None:
+    """``has_dropout_layers`` returns True when Dropout is present, else False."""
     assert has_dropout_layers(DummyCNNModel(with_dropout=True)) is True
     assert has_dropout_layers(DummyCNNModel(with_dropout=False)) is False
 
 
 def test_inject_dropout_before_classifier_wraps_linear() -> None:
+    """Injection wraps a bare ``nn.Linear`` classifier in a Dropout -> Linear stack."""
     model = DummyCNNModel(with_dropout=False)
     model = inject_dropout_before_classifier(model, p=0.3)
 
@@ -74,6 +100,7 @@ def test_inject_dropout_before_classifier_wraps_linear() -> None:
 
 
 def test_inject_dropout_before_classifier_is_idempotent() -> None:
+    """Re-injecting on a model that already has a leading Dropout is a no-op."""
     model = DummyCNNModel(with_dropout=False)
     model = inject_dropout_before_classifier(model, p=0.3)
     n_children_before = len(list(model.classifier.children()))
@@ -85,6 +112,7 @@ def test_inject_dropout_before_classifier_is_idempotent() -> None:
 
 
 def test_mc_dropout_mode_activates_only_dropout() -> None:
+    """Inside the context only Dropout is in train mode; BatchNorm stays in eval."""
     model = DummyCNNModel(with_dropout=True)
     model.eval()  # everything off to start
 
@@ -98,6 +126,7 @@ def test_mc_dropout_mode_activates_only_dropout() -> None:
 
 
 def test_mc_dropout_mode_restores_previous_training_state() -> None:
+    """The previous module training state is restored on context exit."""
     model = DummyCNNModel(with_dropout=True)
     model.train()  # whole model was in train mode before entering
 
@@ -112,6 +141,7 @@ def test_mc_dropout_mode_restores_previous_training_state() -> None:
 # uncertainty.py
 # ----------------------------------------------------------------------------
 def test_decompose_uncertainty_shapes_and_keys() -> None:
+    """``decompose_uncertainty`` returns expected keys and shapes; total = ep + al."""
     torch.manual_seed(0)
     t, batch, n_classes = 20, 5, 3
     logits = torch.randn(t, batch, n_classes)
@@ -130,7 +160,9 @@ def test_decompose_uncertainty_shapes_and_keys() -> None:
 def test_decompose_uncertainty_zero_variance_gives_zero_epistemic() -> None:
     """If every MC sample is identical, epistemic uncertainty must be exactly 0."""
     single_pass = torch.tensor([[0.7, 0.2, 0.1]])
-    mc_probs = single_pass.unsqueeze(0).repeat(10, 1, 1)  # [T=10, B=1, C=3], all identical
+    mc_probs = single_pass.unsqueeze(0).repeat(
+        10, 1, 1
+    )  # [T=10, B=1, C=3], all identical
 
     stats = decompose_uncertainty(mc_probs, class_dim=-1)
 
@@ -142,17 +174,14 @@ def test_decompose_uncertainty_zero_variance_gives_zero_epistemic() -> None:
 # bayesian_wrapper.py
 # ----------------------------------------------------------------------------
 def test_bayesian_wrapper_raises_without_dropout() -> None:
+    """Wrapping a Dropout-less model with ``inject_dropout=False`` raises ValueError."""
     model = DummyCNNModel(with_dropout=False)
-    try:
+    with pytest.raises(ValueError, match="Dropout"):
         BayesianModelWrapper(base_model=model, n_samples=5, inject_dropout=False)
-    except ValueError as exc:
-        assert "Dropout" in str(exc)
-    else:
-        msg = "Expected ValueError when wrapping a model with no Dropout layers."
-        raise AssertionError(msg)
 
 
 def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
+    """``infer_batch`` returns mean probs and appends a per-batch uncertainty dict."""
     model = DummyCNNModel(with_dropout=True)
     model.eval()
 
@@ -169,4 +198,3 @@ def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
     for key in ("mean_probs", "aleatoric", "epistemic", "total", "mc_probs"):
         assert key in batch_stats
     assert batch_stats["mc_probs"].shape == (8, 4, 2)  # [T, B, C]
-
diff --git a/tiatoolbox/models/architecture/bayesian_wrapper.py b/tiatoolbox/models/architecture/bayesian_wrapper.py
index c84a29c5bc..d6adb7d406 100644
--- a/tiatoolbox/models/architecture/bayesian_wrapper.py
+++ b/tiatoolbox/models/architecture/bayesian_wrapper.py
@@ -12,9 +12,8 @@
 
 from __future__ import annotations
 
-from typing import Any
+from typing import TYPE_CHECKING
 
-import numpy as np
 import torch
 
 from tiatoolbox.models.engine.mc_dropout import (
@@ -24,6 +23,9 @@
 )
 from tiatoolbox.models.engine.uncertainty import decompose_uncertainty
 
+if TYPE_CHECKING:  # pragma: no cover
+    import numpy as np
+
 
 class BayesianModelWrapper(torch.nn.Module):
     """Wrap a TIAToolbox model and add Monte Carlo Dropout at `infer_batch`.
@@ -45,9 +47,33 @@ def __init__(
         base_model: torch.nn.Module,
         n_samples: int = 30,
         class_dim: int = -1,
+        *,
         inject_dropout: bool = True,
         dropout_p: float = 0.2,
     ) -> None:
+        """Initialize :class:`BayesianModelWrapper`.
+
+        Args:
+            base_model (torch.nn.Module):
+                TIAToolbox model to wrap. Must expose an ``infer_batch`` static
+                method matching ``(model, images, device)``.
+            n_samples (int):
+                Number of Monte Carlo forward passes per batch.
+            class_dim (int):
+                Class axis in the per-sample probabilities (excluding T).
+            inject_dropout (bool):
+                If True and the model has no Dropout layers, insert one before
+                the linear classifier automatically.
+            dropout_p (float):
+                Dropout probability used by ``inject_dropout``.
+
+        Raises:
+            ValueError:
+                If the model has no Dropout layers and ``inject_dropout`` is False.
+            RuntimeError:
+                If dropout injection was requested but failed.
+
+        """
         super().__init__()
         if not has_dropout_layers(base_model):
             if not inject_dropout:
@@ -72,11 +98,28 @@ def __init__(
         # `coordinates` / `probabilities` outputs.
         self.uncertainty_stats: list[dict[str, np.ndarray]] = []
 
-    def __getattr__(self, name: str) -> Any:
-        # Delegate undefined attributes to the wrapped model (e.g. preproc_func,
-        # postproc_func, class_dict) so the engine finds them via _get_model_attr.
-        # `__getattr__` is invoked only when normal lookup fails, so by the time
-        # we reach here `self.base_model` is already accessible normally.
+    def __getattr__(self, name: str) -> object:
+        """Delegate undefined attributes to the wrapped model.
+
+        This lets the engine find attributes such as ``preproc_func``,
+        ``postproc_func``, and ``class_dict`` on the wrapped model through
+        ``_get_model_attr``. ``__getattr__`` is invoked only when normal
+        lookup fails, so by the time we reach here ``self.base_model`` is
+        already accessible normally.
+
+        Args:
+            name (str):
+                Attribute name to look up.
+
+        Returns:
+            object:
+                The attribute found on the wrapped model.
+
+        Raises:
+            AttributeError:
+                If the attribute is not found on the wrapper or the model.
+
+        """
         try:
             return super().__getattr__(name)
         except AttributeError:
@@ -89,7 +132,7 @@ def __getattr__(self, name: str) -> Any:
 
     @staticmethod
     def infer_batch(
-        model: "BayesianModelWrapper",
+        model: BayesianModelWrapper,
         batch_data: torch.Tensor,
         device: str = "cpu",
     ) -> np.ndarray:
@@ -120,6 +163,9 @@ def infer_batch(
             }
         )
 
-        return stats["mean_probs"].to(
-            next(wrapper.base_model.parameters()).device
-        ).cpu().numpy()
+        return (
+            stats["mean_probs"]
+            .to(next(wrapper.base_model.parameters()).device)
+            .cpu()
+            .numpy()
+        )
diff --git a/tiatoolbox/models/engine/mc_dropout.py b/tiatoolbox/models/engine/mc_dropout.py
index 8c2c35a52f..ac8b0dcdb6 100644
--- a/tiatoolbox/models/engine/mc_dropout.py
+++ b/tiatoolbox/models/engine/mc_dropout.py
@@ -2,8 +2,14 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from torch import nn
 
+if TYPE_CHECKING:  # pragma: no cover
+    from types import TracebackType
+    from typing import Self
+
 _DROPOUT_TYPES = (
     nn.Dropout,
     nn.Dropout1d,
@@ -13,18 +19,38 @@
     nn.FeatureAlphaDropout,
 )
 
+# Minimum number of children a ``nn.Sequential`` classifier must have before we
+# consider its leading ``nn.Dropout`` as already injected.
+_MIN_SEQ_LEN_FOR_DROPOUT = 2
+
 
-class mc_dropout_mode:
+class mc_dropout_mode:  # noqa: N801 - intentional snake_case context-manager name
     """Context manager: enable only Dropout layers, everything else stays in eval."""
 
     def __init__(self, model: nn.Module) -> None:
+        """Initialize :class:`mc_dropout_mode`.
+
+        Args:
+            model (nn.Module):
+                The model whose Dropout layers will be put in train mode
+                inside the context.
+
+        """
         self.model = model
         # Save training state of ALL modules (not just Dropout) so BatchNorm
         # and the root module are restored correctly on exit, since
         # `self.model.eval()` puts the whole tree in eval mode.
         self._prev_states: dict[nn.Module, bool] = {}
 
-    def __enter__(self) -> nn.Module:
+    def __enter__(self) -> Self:
+        """Enter the context, switching only Dropout layers to train mode.
+
+        Returns:
+            Self:
+                The model being managed (returned so callers can use it inside
+                the ``with`` block).
+
+        """
         self._prev_states = {m: m.training for m in self.model.modules()}
 
         self.model.eval()
@@ -34,15 +60,39 @@ def __enter__(self) -> nn.Module:
 
         return self.model
 
-    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit the context, restoring each module's prior training state.
+
+        Args:
+            exc_type (type[BaseException] | None):
+                Type of the raised exception, if any.
+            exc_val (BaseException | None):
+                The raised exception, if any.
+            exc_tb (TracebackType | None):
+                The associated traceback, if any.
+
+        """
         for module, was_training in self._prev_states.items():
             module.train(was_training)
 
-        return False  # do not suppress exceptions
-
 
 def has_dropout_layers(model: nn.Module) -> bool:
-    """Check if the model has any dropout layers."""
+    """Check if the model has any dropout layers.
+
+    Args:
+        model (nn.Module):
+            The model to inspect.
+
+    Returns:
+        bool:
+            True if at least one submodule is an instance of a Dropout type.
+
+    """
     return any(isinstance(m, _DROPOUT_TYPES) for m in model.modules())
 
 
@@ -68,7 +118,11 @@ def inject_dropout_before_classifier(
             f"got {type(model).__name__}."
         )
         raise TypeError(msg)
-    if isinstance(head, nn.Sequential) and len(head) >= 2 and isinstance(head[0], nn.Dropout):
+    if (
+        isinstance(head, nn.Sequential)
+        and len(head) >= _MIN_SEQ_LEN_FOR_DROPOUT
+        and isinstance(head[0], nn.Dropout)
+    ):
         return model
     if not isinstance(head, nn.Linear):
         msg = (
diff --git a/tiatoolbox/models/engine/uncertainty.py b/tiatoolbox/models/engine/uncertainty.py
index 900635cd92..6cd74e847f 100644
--- a/tiatoolbox/models/engine/uncertainty.py
+++ b/tiatoolbox/models/engine/uncertainty.py
@@ -1,11 +1,14 @@
-"""
-Uncertainty Quantification with Monte Carlo (MC Dropout).
+"""Uncertainty Quantification with Monte Carlo (MC Dropout).
+
 Gal & Ghahramani (2015), used by Policastro (2020) WSI analysis project.
 """
 
 from __future__ import annotations
 
-import torch
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:  # pragma: no cover
+    import torch
 
 
 def decompose_uncertainty(
@@ -20,9 +23,10 @@ def decompose_uncertainty(
 
     Returns:
         dict: mean_probs, epistemic, aleatoric, total uncertainty.
+
     """
-    mean_probs = mc_probs.mean(dim=0)          # [..., C]
-    var_probs = mc_probs.var(dim=0)            # [..., C]
+    mean_probs = mc_probs.mean(dim=0)  # [..., C]
+    var_probs = mc_probs.var(dim=0)  # [..., C]
 
     epistemic = var_probs.mean(dim=class_dim)  # variance across T passes
     aleatoric = (mean_probs * (1 - mean_probs)).mean(dim=class_dim)
@@ -33,4 +37,4 @@ def decompose_uncertainty(
         "epistemic": epistemic,
         "aleatoric": aleatoric,
         "total": total,
-    }
\ No newline at end of file
+    }

From 0ce787f20d1ecdeb84fffd5ee97eda8085206ae3 Mon Sep 17 00:00:00 2001
From: Piero <piero@example.com>
Date: Mon, 17 Aug 2026 11:33:11 +0200
Subject: [PATCH 3/6] test: add missing unit tests for MC-Dropout coverage

- bayesian_wrapper: cover RuntimeError on failed dropout injection (lines 85-91)
  and __getattr__ delegation/AttributeError (lines 125-131)
- mc_dropout: cover TypeError branches in inject_dropout_before_classifier,
  AlphaDropout/FeatureAlphaDropout detection, and mc_dropout_mode.__exit__
  with exception restore

Fixes Codecov coverage gaps reported in PR #1104
---
 tests/test_mc_dropout_units.py | 77 +++++++++++++++++++++++++++++++++-
 1 file changed, 76 insertions(+), 1 deletion(-)

diff --git a/tests/test_mc_dropout_units.py b/tests/test_mc_dropout_units.py
index a0fe372c5f..802ac41077 100644
--- a/tests/test_mc_dropout_units.py
+++ b/tests/test_mc_dropout_units.py
@@ -137,6 +137,48 @@ def test_mc_dropout_mode_restores_previous_training_state() -> None:
     assert model.bn.training is True
 
 
+def test_inject_dropout_before_classifier_raises_typeerror_no_classifier() -> None:
+    """TypeError if model has no `classifier` attribute."""
+    model = nn.Sequential(nn.Linear(10, 2))
+    with pytest.raises(TypeError, match="classifier"):
+        inject_dropout_before_classifier(model)
+
+
+def test_inject_dropout_before_classifier_raises_typeerror_invalid_head() -> None:
+    """TypeError if classifier is neither Linear nor Sequential with leading Dropout."""
+    model = DummyCNNModel(with_dropout=False)
+    model.classifier = nn.Conv2d(4, 2, 1)
+    with pytest.raises(TypeError, match="nn.Linear"):
+        inject_dropout_before_classifier(model)
+
+
+def test_has_dropout_layers_detects_alpha_dropout() -> None:
+    """has_dropout_layers detects AlphaDropout and FeatureAlphaDropout."""
+    model_alpha = nn.Sequential(nn.Linear(10, 2), nn.AlphaDropout())
+    model_feat_alpha = nn.Sequential(nn.Linear(10, 2), nn.FeatureAlphaDropout())
+    assert has_dropout_layers(model_alpha) is True
+    assert has_dropout_layers(model_feat_alpha) is True
+
+
+def test_mc_dropout_mode_exit_restores_on_exception() -> None:
+    """__exit__ restores training state even when an exception is raised."""
+    model = DummyCNNModel(with_dropout=True)
+    model.eval()
+
+    class TestError(Exception):
+        pass
+
+    try:
+        with mc_dropout_mode(model):
+            raise TestError("boom")
+    except TestError:
+        pass
+
+    assert model.training is False
+    assert model.bn.training is False
+    assert model.classifier[0].training is False
+
+
 # ----------------------------------------------------------------------------
 # uncertainty.py
 # ----------------------------------------------------------------------------
@@ -180,6 +222,39 @@ def test_bayesian_wrapper_raises_without_dropout() -> None:
         BayesianModelWrapper(base_model=model, n_samples=5, inject_dropout=False)
 
 
+def test_bayesian_wrapper_raises_runtime_error_if_injection_fails() -> None:
+    """RuntimeError se inject_dropout_before_classifier non aggiunge dropout."""
+    from unittest.mock import patch
+
+    model_without_dropout = DummyCNNModel(with_dropout=False)
+
+    with patch(
+        "tiatoolbox.models.architecture.bayesian_wrapper.inject_dropout_before_classifier",
+        return_value=model_without_dropout,
+    ):
+        with pytest.raises(RuntimeError, match="failed to insert"):
+            BayesianModelWrapper(model_without_dropout, inject_dropout=True)
+
+
+def test_getattr_delegates_to_base_model() -> None:
+    """__getattr__ deve delegare al base_model per attributi non definiti sul wrapper."""
+    base_model = DummyCNNModel(with_dropout=True)
+    base_model.preproc_func = lambda x: x
+
+    wrapper = BayesianModelWrapper(base_model, inject_dropout=False)
+
+    assert wrapper.preproc_func is base_model.preproc_func
+
+
+def test_getattr_raises_if_attribute_missing_everywhere() -> None:
+    """__getattr__ deve sollevare AttributeError se l'attributo non esiste da nessuna parte."""
+    base_model = DummyCNNModel(with_dropout=True)
+    wrapper = BayesianModelWrapper(base_model, inject_dropout=False)
+
+    with pytest.raises(AttributeError, match="has no attribute"):
+        _ = wrapper.this_attribute_does_not_exist_anywhere
+
+
 def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
     """``infer_batch`` returns mean probs and appends a per-batch uncertainty dict."""
     model = DummyCNNModel(with_dropout=True)
@@ -197,4 +272,4 @@ def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
     batch_stats = wrapper.uncertainty_stats[0]
     for key in ("mean_probs", "aleatoric", "epistemic", "total", "mc_probs"):
         assert key in batch_stats
-    assert batch_stats["mc_probs"].shape == (8, 4, 2)  # [T, B, C]
+    assert batch_stats["mc_probs"].shape == (8, 4, 2)  # [T, B, C]
\ No newline at end of file

From af16e94e70ae5537a7daf655ea7490e27d4e97d3 Mon Sep 17 00:00:00 2001
From: Piero <piero@example.com>
Date: Tue, 18 Aug 2026 01:31:28 +0200
Subject: [PATCH 4/6] test+fix: add missing unit tests and fix lint issues for
 MC-Dropout

   - Add tests covering Codecov gaps in bayesian_wrapper (lines 85-91, 125-131)
     and mc_dropout (TypeError branches, AlphaDropout, __exit__ with exception)
   - Fix FLK-D402: rewrite DummyCNNModel.infer_batch docstring
   - Fix PTC-W0063: wrap next() in try-except with CPU fallback
---
 tests/test_mc_dropout_units.py                | 58 +++++++++++--------
 .../models/architecture/bayesian_wrapper.py   | 12 ++--
 2 files changed, 40 insertions(+), 30 deletions(-)

diff --git a/tests/test_mc_dropout_units.py b/tests/test_mc_dropout_units.py
index 802ac41077..6c4dea741d 100644
--- a/tests/test_mc_dropout_units.py
+++ b/tests/test_mc_dropout_units.py
@@ -14,6 +14,8 @@
 
 from __future__ import annotations
 
+from unittest.mock import patch
+
 import numpy as np
 import pytest
 import torch
@@ -73,7 +75,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
     def infer_batch(
         model: DummyCNNModel, batch_data: torch.Tensor, device: str = "cpu"
     ) -> torch.Tensor:
-        """Mimic tiatoolbox's CNNModel.infer_batch(model, images, device) signature."""
+        """Run a forward pass and return softmax probabilities.
+
+        Mimics tiatoolbox's CNNModel.infer_batch(model, images, device) interface.
+        """
         with torch.no_grad():
             logits = model(batch_data.to(device))
             return torch.softmax(logits, dim=-1)
@@ -140,7 +145,7 @@ def test_mc_dropout_mode_restores_previous_training_state() -> None:
 def test_inject_dropout_before_classifier_raises_typeerror_no_classifier() -> None:
     """TypeError if model has no `classifier` attribute."""
     model = nn.Sequential(nn.Linear(10, 2))
-    with pytest.raises(TypeError, match="classifier"):
+    with pytest.raises(TypeError, match=r"classifier"):
         inject_dropout_before_classifier(model)
 
 
@@ -148,7 +153,7 @@ def test_inject_dropout_before_classifier_raises_typeerror_invalid_head() -> Non
     """TypeError if classifier is neither Linear nor Sequential with leading Dropout."""
     model = DummyCNNModel(with_dropout=False)
     model.classifier = nn.Conv2d(4, 2, 1)
-    with pytest.raises(TypeError, match="nn.Linear"):
+    with pytest.raises(TypeError, match=r"nn\.Linear"):
         inject_dropout_before_classifier(model)
 
 
@@ -160,19 +165,23 @@ def test_has_dropout_layers_detects_alpha_dropout() -> None:
     assert has_dropout_layers(model_feat_alpha) is True
 
 
+class _TestError(Exception):
+    """Test exception for context manager testing."""
+
+
+def _raise_test_error() -> None:
+    """Raise a test exception for context manager testing."""
+    msg = "boom"
+    raise _TestError(msg)
+
+
 def test_mc_dropout_mode_exit_restores_on_exception() -> None:
     """__exit__ restores training state even when an exception is raised."""
     model = DummyCNNModel(with_dropout=True)
     model.eval()
 
-    class TestError(Exception):
-        pass
-
-    try:
-        with mc_dropout_mode(model):
-            raise TestError("boom")
-    except TestError:
-        pass
+    with pytest.raises(_TestError), mc_dropout_mode(model):
+        _raise_test_error()
 
     assert model.training is False
     assert model.bn.training is False
@@ -218,26 +227,27 @@ def test_decompose_uncertainty_zero_variance_gives_zero_epistemic() -> None:
 def test_bayesian_wrapper_raises_without_dropout() -> None:
     """Wrapping a Dropout-less model with ``inject_dropout=False`` raises ValueError."""
     model = DummyCNNModel(with_dropout=False)
-    with pytest.raises(ValueError, match="Dropout"):
+    with pytest.raises(ValueError, match=r"Dropout"):
         BayesianModelWrapper(base_model=model, n_samples=5, inject_dropout=False)
 
 
 def test_bayesian_wrapper_raises_runtime_error_if_injection_fails() -> None:
-    """RuntimeError se inject_dropout_before_classifier non aggiunge dropout."""
-    from unittest.mock import patch
-
+    """RuntimeError if inject_dropout_before_classifier fails to add dropout."""
     model_without_dropout = DummyCNNModel(with_dropout=False)
 
-    with patch(
-        "tiatoolbox.models.architecture.bayesian_wrapper.inject_dropout_before_classifier",
-        return_value=model_without_dropout,
+    with (
+        patch(
+            "tiatoolbox.models.architecture.bayesian_wrapper."
+            "inject_dropout_before_classifier",
+            return_value=model_without_dropout,
+        ),
+        pytest.raises(RuntimeError, match=r"failed to insert"),
     ):
-        with pytest.raises(RuntimeError, match="failed to insert"):
-            BayesianModelWrapper(model_without_dropout, inject_dropout=True)
+        BayesianModelWrapper(model_without_dropout, inject_dropout=True)
 
 
 def test_getattr_delegates_to_base_model() -> None:
-    """__getattr__ deve delegare al base_model per attributi non definiti sul wrapper."""
+    """__getattr__ delegates to base_model for attributes not on wrapper."""
     base_model = DummyCNNModel(with_dropout=True)
     base_model.preproc_func = lambda x: x
 
@@ -247,11 +257,11 @@ def test_getattr_delegates_to_base_model() -> None:
 
 
 def test_getattr_raises_if_attribute_missing_everywhere() -> None:
-    """__getattr__ deve sollevare AttributeError se l'attributo non esiste da nessuna parte."""
+    """__getattr__ raises AttributeError if attribute missing everywhere."""
     base_model = DummyCNNModel(with_dropout=True)
     wrapper = BayesianModelWrapper(base_model, inject_dropout=False)
 
-    with pytest.raises(AttributeError, match="has no attribute"):
+    with pytest.raises(AttributeError, match=r"has no attribute"):
         _ = wrapper.this_attribute_does_not_exist_anywhere
 
 
@@ -272,4 +282,4 @@ def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
     batch_stats = wrapper.uncertainty_stats[0]
     for key in ("mean_probs", "aleatoric", "epistemic", "total", "mc_probs"):
         assert key in batch_stats
-    assert batch_stats["mc_probs"].shape == (8, 4, 2)  # [T, B, C]
\ No newline at end of file
+    assert batch_stats["mc_probs"].shape == (8, 4, 2)  # [T, B, C]
diff --git a/tiatoolbox/models/architecture/bayesian_wrapper.py b/tiatoolbox/models/architecture/bayesian_wrapper.py
index d6adb7d406..472bca28d2 100644
--- a/tiatoolbox/models/architecture/bayesian_wrapper.py
+++ b/tiatoolbox/models/architecture/bayesian_wrapper.py
@@ -163,9 +163,9 @@ def infer_batch(
             }
         )
 
-        return (
-            stats["mean_probs"]
-            .to(next(wrapper.base_model.parameters()).device)
-            .cpu()
-            .numpy()
-        )
+        try:
+            target_device = next(wrapper.base_model.parameters()).device
+        except StopIteration:
+            target_device = torch.device("cpu")
+
+        return stats["mean_probs"].to(target_device).cpu().numpy()

From f3b5ae99cb4b50f2637be5136364fbf4a1b41374 Mon Sep 17 00:00:00 2001
From: PierWSL2 <piero.policastro@gmail.com>
Date: Tue, 15 Sep 2026 21:09:29 +0200
Subject: [PATCH 5/6] Fix eval, Fix Codecov, Introducing use example

---
 tests/test_mc_dropout_units.py                | 75 +++++++++++++++++++
 .../models/architecture/bayesian_wrapper.py   | 36 +++++++++
 tiatoolbox/models/architecture/vanilla.py     |  6 +-
 .../models/engine/deep_feature_extractor.py   |  5 ++
 tiatoolbox/models/engine/engine_abc.py        |  4 +
 tiatoolbox/models/engine/mc_dropout.py        | 12 ++-
 .../models/engine/multi_task_segmentor.py     |  8 ++
 tiatoolbox/models/engine/prompt_segmentor.py  |  7 +-
 .../models/engine/semantic_segmentor.py       |  6 +-
 tiatoolbox/models/engine/uncertainty.py       |  8 ++
 10 files changed, 161 insertions(+), 6 deletions(-)

diff --git a/tests/test_mc_dropout_units.py b/tests/test_mc_dropout_units.py
index 6c4dea741d..9bc9a90334 100644
--- a/tests/test_mc_dropout_units.py
+++ b/tests/test_mc_dropout_units.py
@@ -265,6 +265,81 @@ def test_getattr_raises_if_attribute_missing_everywhere() -> None:
         _ = wrapper.this_attribute_does_not_exist_anywhere
 
 
+class _ParamlessDropoutModel(nn.Module):
+    """Dropout-only model with no parameters, to hit the StopIteration branch."""
+
+    def __init__(self) -> None:
+        """Initialize :class:`_ParamlessDropoutModel`."""
+        super().__init__()
+        self.dropout = nn.Dropout(0.5)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Pass input through the dropout layer."""
+        return self.dropout(x)
+
+    @staticmethod
+    def infer_batch(
+        model: _ParamlessDropoutModel, batch_data: torch.Tensor, device: str = "cpu"
+    ) -> torch.Tensor:
+        """Return softmax probabilities, mimicking CNNModel.infer_batch."""
+        with torch.no_grad():
+            return torch.softmax(model(batch_data.to(device)), dim=-1)
+
+
+def test_getattr_raises_when_base_model_unset() -> None:
+    """__getattr__ raises AttributeError when base_model has not been assigned."""
+    wrapper = BayesianModelWrapper.__new__(BayesianModelWrapper)
+    nn.Module.__init__(wrapper)  # empty module: no base_model registered
+
+    with pytest.raises(AttributeError, match=r"has no attribute"):
+        _ = wrapper.any_attribute
+
+
+def test_bayesian_wrapper_infer_batch_model_without_parameters() -> None:
+    """infer_batch falls back to CPU device when the model has no parameters."""
+    model = _ParamlessDropoutModel()
+    wrapper = BayesianModelWrapper(base_model=model, n_samples=3, class_dim=-1)
+    dummy_batch = torch.rand(2, 4)
+
+    mean_probs = wrapper.infer_batch(wrapper, dummy_batch, device="cpu")
+
+    assert isinstance(mean_probs, np.ndarray)
+    assert mean_probs.shape == (2, 4)
+    assert len(wrapper.uncertainty_stats) == 1
+
+
+def test_vanilla_infer_batch_preserves_mc_dropout() -> None:
+    """``_infer_batch`` must not revert dropout layers set by ``mc_dropout_mode``.
+
+    Regression test: ``_infer_batch`` used to call ``model.eval()`` internally,
+    which silently deactivated the Dropout layers enabled by
+    :class:`mc_dropout_mode`, making all MC samples identical (epistemic = 0)
+    for every CNNModel/TimmModel-based model.
+    """
+    import importlib  # noqa: PLC0415
+
+    vanilla = importlib.import_module("tiatoolbox.models.architecture.vanilla")
+
+    model = DummyCNNModel(with_dropout=True)
+    model.eval()
+    dummy_batch = torch.rand(4, 8, 8, 3)  # NHWC, as _infer_batch expects
+
+    with mc_dropout_mode(model):
+        outputs = [
+            vanilla._infer_batch(model, dummy_batch, device="cpu") for _ in range(5)
+        ]
+        # Dropout layers must still be in train mode after repeated calls
+        assert model.classifier[0].training is True
+
+    stacks = np.stack(outputs, axis=0)
+    # At least one forward pass must differ (dropout was actually active)
+    assert not np.allclose(stacks, stacks[0])
+
+    # Outside the context the model is back in full eval mode
+    assert model.training is False
+    assert model.classifier[0].training is False
+
+
 def test_bayesian_wrapper_infer_batch_populates_uncertainty_stats() -> None:
     """``infer_batch`` returns mean probs and appends a per-batch uncertainty dict."""
     model = DummyCNNModel(with_dropout=True)
diff --git a/tiatoolbox/models/architecture/bayesian_wrapper.py b/tiatoolbox/models/architecture/bayesian_wrapper.py
index 472bca28d2..ce6a042054 100644
--- a/tiatoolbox/models/architecture/bayesian_wrapper.py
+++ b/tiatoolbox/models/architecture/bayesian_wrapper.py
@@ -40,6 +40,42 @@ class BayesianModelWrapper(torch.nn.Module):
     (mean_probs, aleatoric, epistemic, total, mc_probs) is appended to
     `self.uncertainty_stats` in inference order, so it can be aligned with
     the `coordinates` array written to the engine's output `.zarr`.
+
+    Example:
+        Wrap a pretrained patch classifier and run inference with the
+        patch predictor, collecting uncertainty per patch:
+
+        >>> from tiatoolbox.models.architecture import get_pretrained_model
+        >>> from tiatoolbox.models.engine.patch_predictor import PatchPredictor
+        >>> from tiatoolbox.models.architecture.bayesian_wrapper import (
+        ...     BayesianModelWrapper,
+        ... )
+        >>>
+        >>> model, _ = get_pretrained_model("resnet18-kather100k")
+        >>> bayesian_model = BayesianModelWrapper(
+        ...     base_model=model,
+        ...     n_samples=30,  # number of MC forward passes
+        ...     inject_dropout=True,  # model ships without Dropout
+        ...     dropout_p=0.2,
+        ... )
+        >>>
+        >>> predictor = PatchPredictor(model=bayesian_model)
+        >>> outputs = predictor.run(images=["sample_img.png"], patch_mode=True)
+        >>>
+        >>> # After predict() finishes, concatenate the per-batch stats:
+        >>> import numpy as np
+        >>> epistemic = np.concatenate(
+        ...     [s["epistemic"] for s in bayesian_model.uncertainty_stats], axis=0
+        ... )  # one value per patch, aligned with outputs["coordinates"]
+        >>> total = np.concatenate(
+        ...     [s["total"] for s in bayesian_model.uncertainty_stats], axis=0
+        ... )
+
+    Note:
+        The wrapped model must be a classification model whose
+        ``infer_batch(model, images, device)`` returns per-sample class
+        probabilities of shape ``[B, C]`` (e.g. :class:`CNNModel`,
+        :class:`TimmModel`).
     """
 
     def __init__(
diff --git a/tiatoolbox/models/architecture/vanilla.py b/tiatoolbox/models/architecture/vanilla.py
index 93e1e59b99..f49a39d08a 100644
--- a/tiatoolbox/models/architecture/vanilla.py
+++ b/tiatoolbox/models/architecture/vanilla.py
@@ -245,8 +245,10 @@ def _infer_batch(
     )  # to NCHW
     img_patches_device = img_patches_device.permute(0, 3, 1, 2).contiguous()
 
-    # Inference mode
-    model.eval()
+    # NOTE: model.eval() is intentionally not called here. Inference
+    # workflows (engines) must set eval mode once before inference. This
+    # allows Monte Carlo Dropout to keep Dropout layers active during
+    # inference without eval() silently reverting them inside this function.
     # Do not compute the gradient (not training)
     with torch.inference_mode():
         output = model(img_patches_device)
diff --git a/tiatoolbox/models/engine/deep_feature_extractor.py b/tiatoolbox/models/engine/deep_feature_extractor.py
index 16fd09ec5c..b464b767e6 100644
--- a/tiatoolbox/models/engine/deep_feature_extractor.py
+++ b/tiatoolbox/models/engine/deep_feature_extractor.py
@@ -44,6 +44,7 @@
 
 import dask.array as da
 import psutil
+import torch
 import zarr
 from dask import compute
 from tqdm.auto import tqdm
@@ -225,6 +226,10 @@ def infer_wsi(
         probabilities_zarr, coordinates_zarr = None, None
 
         probabilities_used_percent = 0
+        # Set inference mode once here; infer_batch no longer calls
+        # model.eval() internally (needed for Monte Carlo Dropout wrappers).
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         infer_batch = self._get_model_attr("infer_batch")
         for batch_data in tqdm_loop:
             batch_output = infer_batch(
diff --git a/tiatoolbox/models/engine/engine_abc.py b/tiatoolbox/models/engine/engine_abc.py
index 208a516db8..5312906cb7 100644
--- a/tiatoolbox/models/engine/engine_abc.py
+++ b/tiatoolbox/models/engine/engine_abc.py
@@ -551,6 +551,10 @@ def infer_patches(
             disable=not self.verbose,
         )
 
+        # Set inference mode once here; infer_batch no longer calls
+        # model.eval() internally (needed for Monte Carlo Dropout wrappers).
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         infer_batch = self._get_model_attr("infer_batch")
         for batch_data in tqdm_loop:
             batch_output = infer_batch(
diff --git a/tiatoolbox/models/engine/mc_dropout.py b/tiatoolbox/models/engine/mc_dropout.py
index ac8b0dcdb6..39e7a2a0a9 100644
--- a/tiatoolbox/models/engine/mc_dropout.py
+++ b/tiatoolbox/models/engine/mc_dropout.py
@@ -25,7 +25,17 @@
 
 
 class mc_dropout_mode:  # noqa: N801 - intentional snake_case context-manager name
-    """Context manager: enable only Dropout layers, everything else stays in eval."""
+    """Context manager: enable only Dropout layers, everything else stays in eval.
+
+    Example:
+        Run 3 stochastic forward passes with dropout active while keeping
+        BatchNorm in eval mode:
+
+        >>> model.eval()
+        >>> with mc_dropout_mode(model):
+        ...     samples = [model(images) for _ in range(3)]
+        >>> # on exit, all modules are back to their previous training state
+    """
 
     def __init__(self, model: nn.Module) -> None:
         """Initialize :class:`mc_dropout_mode`.
diff --git a/tiatoolbox/models/engine/multi_task_segmentor.py b/tiatoolbox/models/engine/multi_task_segmentor.py
index fec5aebe9b..d59f021b51 100644
--- a/tiatoolbox/models/engine/multi_task_segmentor.py
+++ b/tiatoolbox/models/engine/multi_task_segmentor.py
@@ -416,6 +416,10 @@ def infer_patches(
         keys = ["probabilities"]
         coordinates = []
 
+        # Set inference mode once here; infer_batch no longer calls
+        # model.eval() internally (needed for Monte Carlo Dropout wrappers).
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         # Expected number of outputs from the model
         infer_batch = self._get_model_attr("infer_batch")
         batch_output = infer_batch(
@@ -591,6 +595,10 @@ def infer_wsi(
             disable=not self.verbose,
         )
 
+        # Set inference mode once here; infer_batch no longer calls
+        # model.eval() internally (needed for Monte Carlo Dropout wrappers).
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         # Expected number of outputs from the model
         infer_batch = self._get_model_attr("infer_batch")
         batch_output = infer_batch(
diff --git a/tiatoolbox/models/engine/prompt_segmentor.py b/tiatoolbox/models/engine/prompt_segmentor.py
index 608251c7c9..1f250c0175 100644
--- a/tiatoolbox/models/engine/prompt_segmentor.py
+++ b/tiatoolbox/models/engine/prompt_segmentor.py
@@ -6,13 +6,12 @@
 from typing import TYPE_CHECKING
 
 import numpy as np
+import torch
 
 from tiatoolbox.models.architecture.sam import SAM
 from tiatoolbox.utils.misc import dict_to_store_semantic_segmentor
 
 if TYPE_CHECKING:  # pragma: no cover
-    import torch
-
     from tiatoolbox.type_hints import IntPair
 
 
@@ -71,6 +70,10 @@ def run(  # skipcq: PYL-W0221
 
         """
         paths = []
+        # Set inference mode once here; infer_batch no longer calls
+        # model.eval() internally (needed for Monte Carlo Dropout wrappers).
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         masks, _ = self.model.infer_batch(
             self.model,
             images,
diff --git a/tiatoolbox/models/engine/semantic_segmentor.py b/tiatoolbox/models/engine/semantic_segmentor.py
index f098a6f755..b929e4f0dd 100644
--- a/tiatoolbox/models/engine/semantic_segmentor.py
+++ b/tiatoolbox/models/engine/semantic_segmentor.py
@@ -389,7 +389,7 @@ def get_dataloader(
             wsireader_kwargs=wsireader_kwargs,
         )
 
-    def infer_wsi(
+    def infer_wsi(  # noqa: PLR0915  (eval-mode setup nudged statement count)
         self: SemanticSegmentor,
         dataloader: DataLoader,
         save_path: Path,
@@ -502,6 +502,10 @@ def infer_wsi(
         min_x, min_y, _, _ = self.mask_padding
         _, canvas_width = masked_output_shape
 
+        # Set inference mode once here; infer_batch no longer calls
+        # model.eval() internally (needed for Monte Carlo Dropout wrappers).
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         infer_batch = self._get_model_attr("infer_batch")
         for batch_idx, batch_data in enumerate(tqdm_loop):
             batch_output = infer_batch(
diff --git a/tiatoolbox/models/engine/uncertainty.py b/tiatoolbox/models/engine/uncertainty.py
index 6cd74e847f..69e6c4f882 100644
--- a/tiatoolbox/models/engine/uncertainty.py
+++ b/tiatoolbox/models/engine/uncertainty.py
@@ -1,6 +1,8 @@
 """Uncertainty Quantification with Monte Carlo (MC Dropout).
 
 Gal & Ghahramani (2015), used by Policastro (2020) WSI analysis project.
+The main entry point is :class:`BayesianModelWrapper
+<tiatoolbox.models.architecture.bayesian_wrapper.BayesianModelWrapper>`.
 """
 
 from __future__ import annotations
@@ -24,6 +26,12 @@ def decompose_uncertainty(
     Returns:
         dict: mean_probs, epistemic, aleatoric, total uncertainty.
 
+    Example:
+        >>> # mc_probs: [T=30 MC passes, B=4 patches, C=2 classes]
+        >>> stats = decompose_uncertainty(mc_probs, class_dim=-1)
+        >>> stats["mean_probs"].shape  # [B, C] predictive probabilities
+        >>> stats["epistemic"].shape  # [B] model (reducible) uncertainty
+        >>> stats["aleatoric"].shape  # [B] data (irreducible) uncertainty
     """
     mean_probs = mc_probs.mean(dim=0)  # [..., C]
     var_probs = mc_probs.var(dim=0)  # [..., C]

From 23dd4c3ae67c1a6c5a7be76e57b42d1bb73cc8de Mon Sep 17 00:00:00 2001
From: PierWSL2 <piero.policastro@gmail.com>
Date: Sat, 19 Sep 2026 12:11:05 +0200
Subject: [PATCH 6/6] Fix for deepsource test

---
 tests/test_mc_dropout_units.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_mc_dropout_units.py b/tests/test_mc_dropout_units.py
index 9bc9a90334..c54804fc56 100644
--- a/tests/test_mc_dropout_units.py
+++ b/tests/test_mc_dropout_units.py
@@ -288,7 +288,7 @@ def infer_batch(
 
 def test_getattr_raises_when_base_model_unset() -> None:
     """__getattr__ raises AttributeError when base_model has not been assigned."""
-    wrapper = BayesianModelWrapper.__new__(BayesianModelWrapper)
+    wrapper = object.__new__(BayesianModelWrapper)
     nn.Module.__init__(wrapper)  # empty module: no base_model registered
 
     with pytest.raises(AttributeError, match=r"has no attribute"):