📝 Update 05-patch-prediction notebook for the New API (#977)

## Summary

This PR updates the patch‑prediction example to align with the new `PatchPredictor` engine and fixes a long‑standing issue in `EngineABC` related to model‑attribute retrieval when using `DataParallel`.

---

## What’s Changed

### 🔧 Example Notebook Updates
- Updated **`examples/05-patch-prediction.ipynb`** to use the new `PatchPredictor` engine API.
- Added a new **“Visualize in TIAViz”** section, allowing readers to directly inspect prediction results inside **TIAViz** for a smoother, more interactive workflow.

### 🐛 EngineABC Bug Fix
- Fixed a bug in **`EngineABC`** where model attributes were incorrectly retrieved from a `DataParallel` wrapper.
- Introduced `_get_model_attr()` to safely unwrap the underlying model when needed.
- This resolves multi‑GPU crashes caused by attributes living on the wrapped module instead of the actual model.

---

## Why This Matters
- Ensures the patch‑prediction example stays up‑to‑date with the latest engine design.
- Improves multi‑GPU stability and prevents confusing attribute‑access errors.
- Enhances the user experience by integrating TIAViz visualization directly into the example workflow.

---

## Testing
- Verified that the updated notebook runs end‑to‑end with the new engine.
- Confirmed that multi‑GPU training and inference no longer crash when accessing model attributes.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Shan E Ahmed Raza <13048456+shaneahmed@users.noreply.github.com>

Author: 3 people

examples/05-patch-prediction.ipynb

@@ -50,7 +50,8 @@ def main(files: list[Path]) -> None:
 
     """
     for path in files:
-        notebook = json.loads(path.read_text())
+        with Path.open(path, encoding="utf-8", errors="ignore") as f:
+            notebook = json.load(f)
         formatted_notebook = format_notebook(copy.deepcopy(notebook))
         changed = any(
             cell != formatted_cell

pre-commit/notebook_markdown_format.py

@@ -10,6 +10,7 @@ flask-cors>=4.0.0
 glymur>=0.12.7
 huggingface_hub>=0.33.3
 imagecodecs>=2022.9.26
+ipywidgets>=8.1.7
 joblib>=1.1.1
 jupyterlab>=3.5.2
 matplotlib>=3.6.2

requirements/requirements.txt

@@ -273,6 +273,13 @@ def test_engine_initalization() -> NoReturn:
     eng = TestEngineABC(model=model, weights=weights_path)
     assert isinstance(eng, EngineABC)
 
+    with pytest.raises(AttributeError):
+        _ = eng._get_model_attr("test_attr")
+
+    model.test_attr = True
+    eng = TestEngineABC(model=model, weights=weights_path)
+    assert eng._get_model_attr("test_attr") is True
+
 
 def test_engine_run() -> NoReturn:
     """Test engine run."""

tests/engines/test_engine_abc.py

@@ -302,8 +302,9 @@ def infer_wsi(
         probabilities_zarr, coordinates_zarr = None, None
 
         probabilities_used_percent = 0
+        infer_batch = self._get_model_attr("infer_batch")
         for batch_data in tqdm_loop:
-            batch_output = self.model.infer_batch(
+            batch_output = infer_batch(
                 self.model,
                 batch_data["image"],
                 device=self.device,

tiatoolbox/models/engine/deep_feature_extractor.py

@@ -65,6 +65,7 @@
 
 if TYPE_CHECKING:  # pragma: no cover
     import os
+    from collections.abc import Callable
 
     from torch.utils.data import DataLoader
 
@@ -375,6 +376,14 @@ def _initialize_model_ioconfig(
 
         return model, None
 
+    def _get_model_attr(self: EngineABC, attr_name: str) -> Callable:
+        """Return a model attribute, unwrapping DataParallel if required."""
+        try:
+            return getattr(self.model, attr_name)
+        except AttributeError:
+            module = getattr(self.model, "module", None)
+            return getattr(module, attr_name)
+
     def get_dataloader(
         self: EngineABC,
         images: str | Path | list[str | Path] | np.ndarray,
@@ -428,7 +437,7 @@ def get_dataloader(
                 auto_get_mask=auto_get_mask,
             )
 
-            dataset.preproc_func = self.model.preproc_func
+            dataset.preproc_func = self._get_model_attr("preproc_func")
 
             # preprocessing must be defined with the dataset
             return torch.utils.data.DataLoader(
@@ -444,7 +453,7 @@ def get_dataloader(
             inputs=images, labels=labels, patch_input_shape=ioconfig.patch_input_shape
         )
 
-        dataset.preproc_func = self.model.preproc_func
+        dataset.preproc_func = self._get_model_attr("preproc_func")
 
         # preprocessing must be defined with the dataset
         return torch.utils.data.DataLoader(
@@ -529,8 +538,9 @@ def infer_patches(
             else self.dataloader
         )
 
+        infer_batch = self._get_model_attr("infer_batch")
         for batch_data in tqdm_loop:
-            batch_output = self.model.infer_batch(
+            batch_output = infer_batch(
                 self.model,
                 batch_data["image"],
                 device=self.device,

tiatoolbox/models/engine/engine_abc.py

@@ -372,7 +372,8 @@ def post_process_patches(
             dict[str, da.Array]: Post-processed predictions as a Dask array.
 
         """
-        predictions = self.model.postproc_func(raw_predictions["probabilities"])
+        postproc_func = self._get_model_attr("postproc_func")
+        predictions = postproc_func(raw_predictions["probabilities"])
         raw_predictions["predictions"] = cast_to_min_dtype(predictions)
         return raw_predictions
 

tiatoolbox/models/engine/patch_predictor.py

@@ -361,7 +361,7 @@ def get_dataloader(
                 auto_get_mask=auto_get_mask,
             )
 
-            dataset.preproc_func = self.model.preproc_func
+            dataset.preproc_func = self._get_model_attr("preproc_func")
             self.output_locations = dataset.outputs
 
             # preprocessing must be defined with the dataset
@@ -477,8 +477,9 @@ def infer_wsi(
             else dataloader.dataset.outputs
         )
 
+        infer_batch = self._get_model_attr("infer_batch")
         for batch_idx, batch_data in enumerate(tqdm_loop):
-            batch_output = self.model.infer_batch(
+            batch_output = infer_batch(
                 self.model,
                 batch_data["image"],
                 device=self.device,