📝 Update 05-patch-prediction notebook for the New API

pre-commit/notebook_markdown_format.py

@@ -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

requirements/requirements.txt

@@ -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

tests/engines/test_engine_abc.py

@@ -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."""

tiatoolbox/models/engine/deep_feature_extractor.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/engine_abc.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/patch_predictor.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/semantic_segmentor.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,