⚡ Improve merge performance

Author: shaneahmed

tiatoolbox/models/dataset/dataset_abc.py

@@ -523,6 +523,10 @@ def __init__(  # skipcq: PY-R1000
     def __getitem__(self: WSIPatchDataset, idx: int) -> dict:
         """Get an item from the dataset."""
         coords = self.inputs[idx]
+        output_locs = None
+        if hasattr(self, "outputs"):
+            output_locs = self.outputs[idx]
+
         # Read image patch from the whole-slide image
         patch = self.reader.read_bounds(
             coords,
@@ -535,6 +539,13 @@ def __getitem__(self: WSIPatchDataset, idx: int) -> dict:
         # Apply preprocessing to selected patch
         patch = self._preproc(patch)
 
+        if output_locs is not None:
+            return {
+                "image": patch,
+                "coords": np.array(coords),
+                "output_locs": output_locs,
+            }
+
         return {"image": patch, "coords": np.array(coords)}
 
 

tiatoolbox/models/engine/engine_abc.py

@@ -13,6 +13,7 @@
 import torch
 import zarr
 from dask import compute, delayed
+from dask.diagnostics import ProgressBar
 from torch import nn
 from typing_extensions import Unpack
 
@@ -669,7 +670,8 @@ def save_predictions(
                     )
                     write_tasks.append(task)
 
-            compute(*write_tasks)
+            with ProgressBar():
+                compute(*write_tasks)
 
             return save_path
 

tiatoolbox/models/engine/semantic_segmentor.py

@@ -24,12 +24,16 @@
 if TYPE_CHECKING:  # pragma: no cover
     import os
 
+    from torch.utils.data import DataLoader
+
     from tiatoolbox.annotation import AnnotationStore
     from tiatoolbox.models.engine.io_config import IOSegmentorConfig
     from tiatoolbox.models.models_abc import ModelABC
     from tiatoolbox.type_hints import Resolution
     from tiatoolbox.wsicore import WSIReader
 
+    from .engine_abc import EngineABC, EngineABCRunParams
+
 
 class SemanticSegmentorRunParams(PredictorRunParams):
     """Class describing the input parameters for the :func:`EngineABC.run()` method.
@@ -364,85 +368,144 @@ def get_dataloader(
             patch_mode=patch_mode,
         )
 
-    def post_process_wsi(
-        self: SemanticSegmentor,
-        raw_predictions: Path,
-        **kwargs: Unpack[PredictorRunParams],
+    def infer_wsi(
+        self: EngineABC,
+        dataloader: DataLoader,
+        **kwargs: EngineABCRunParams,
     ) -> Path:
-        """Returns an array from raw predictions.
+        """Model inference on a WSI.
 
-        Merges raw predictions from individual patches into a single prediction array if
-        patch_mode is False.
+        Args:
+            dataloader (DataLoader):
+                A torch dataloader to process WSIs.
+
+            save_path (Path):
+                Path to save the intermediate output. The intermediate output is saved
+                in a zarr file.
+            **kwargs (EngineABCRunParams):
+                Keyword Args to update setup_patch_dataset() method attributes. See
+                :class:`EngineRunParams` for accepted keyword arguments.
+
+        Returns:
+            save_path (Path):
+                Path to zarr file where intermediate output is saved.
 
         """
         _ = kwargs.get("return_probabilities")
-        progress_bar = None
+
         tqdm = get_tqdm()
 
-        if self.verbose:
-            progress_bar = tqdm(
-                total=len(self.output_locations),
-                leave=False,
-                desc="Merging Patch Outputs",
-            )
+        progress_bar = (
+            tqdm(total=len(dataloader), leave=self.patch_mode, desc="Inferring patches")
+            if self.verbose
+            else None
+        )
 
-        num_post_proc_workers = self.num_post_proc_workers
+        keys = ["coordinates"]
 
-        if num_post_proc_workers is not None and num_post_proc_workers > 0:
-            dask.config.set(scheduler="threads", num_workers=num_post_proc_workers)
-        else:
-            dask.config.set(scheduler="threads")
+        if self.return_labels:
+            keys.append("labels")
 
-        dask_patch_probabilities = raw_predictions["probabilities"]
+        raw_predictions = dict.fromkeys(keys)
 
-        # --- Calculate canvas parameters from Dask array and locations ---
         max_location = np.max(self.output_locations, axis=0)
+
+        out_ = self.model.infer_batch(
+            self.model,
+            torch.from_numpy(dataloader.dataset[0]["image"][None, :, :, :]),
+            device=self.device,
+        )
+
         merged_shape = (
             max_location[3],
             max_location[2],
-            dask_patch_probabilities.shape[3],
+            out_["probabilities"].shape[3],
         )
 
         # creating dask arrays for faster processing
         merged_probabilities = da.zeros(
             shape=merged_shape,
-            dtype=dask_patch_probabilities.dtype,
+            dtype=out_["probabilities"].dtype,
             chunks=merged_shape,
         )
 
         merged_weights = da.zeros(
-            shape=merged_shape,
+            shape=merged_shape[:2],
             dtype=int,
-            chunks=merged_shape,
+            chunks=merged_shape[:2],
         )
 
-        for idx, location in enumerate(self.output_locations):
-            start_x, start_y, end_x, end_y = location
-            patch_probs = dask_patch_probabilities[
-                idx, 0 : end_y - start_y, 0 : end_x - start_x, :
-            ]
-            merged_probabilities[start_y:end_y, start_x:end_x, :] = (
-                merged_probabilities[start_y:end_y, start_x:end_x, :] + patch_probs
+        for _, batch_data in enumerate(dataloader):
+            batch_output = self.model.infer_batch(
+                self.model,
+                batch_data["image"],
+                device=self.device,
             )
-            merged_weights[start_y:end_y, start_x:end_x] = (
-                merged_weights[start_y:end_y, start_x:end_x] + 1
+
+            batch_output["coordinates"] = self._get_coordinates(batch_data)
+
+            if self.return_labels:  # be careful of `s`
+                if isinstance(batch_data["label"], torch.Tensor):
+                    batch_output["labels"] = batch_data["label"].numpy()
+                else:
+                    batch_output["labels"] = np.array(batch_data["label"])
+
+            output_locs = batch_data["output_locs"]
+
+            for idx, location in enumerate(output_locs.numpy()):
+                start_x, start_y, end_x, end_y = location
+                patch_probs = batch_output["probabilities"][
+                    idx, 0 : end_y - start_y, 0 : end_x - start_x, :
+                ]
+                merged_probabilities[start_y:end_y, start_x:end_x, :] = (
+                    merged_probabilities[start_y:end_y, start_x:end_x, :] + patch_probs
+                )
+                merged_weights[start_y:end_y, start_x:end_x] = (
+                    merged_weights[start_y:end_y, start_x:end_x] + 1
+                )
+
+            del batch_output["probabilities"]
+            raw_predictions = self._update_model_output(
+                raw_predictions=raw_predictions,
+                raw_output=batch_output,
             )
+
             if progress_bar:
                 progress_bar.update()
 
+        merged_weights = da.maximum(merged_weights, 1)
+        raw_predictions["probabilities"] = (
+            merged_probabilities / merged_weights[:, :, None]
+        )
+
         if progress_bar:
             progress_bar.close()
 
-        # Normalize where weight > 1
-        final_probabilities_dask = da.where(
-            merged_weights > 1,
-            merged_probabilities / merged_weights,
-            merged_probabilities,
-        )
+        return raw_predictions
+
+    def post_process_wsi(
+        self: SemanticSegmentor,
+        raw_predictions: Path,
+        **kwargs: Unpack[PredictorRunParams],
+    ) -> Path:
+        """Returns an array from raw predictions.
+
+        Merges raw predictions from individual patches into a single prediction array if
+        patch_mode is False.
+
+        """
+        _ = kwargs.get("return_probabilities")
+
+        num_post_proc_workers = self.num_post_proc_workers
+
+        if num_post_proc_workers is not None and num_post_proc_workers > 0:
+            dask.config.set(scheduler="threads", num_workers=num_post_proc_workers)
+        else:
+            dask.config.set(scheduler="threads")
 
         # Applying Post-Processing
         raw_predictions["predictions"] = self.model.postproc_func(
-            final_probabilities_dask,
+            raw_predictions["probabilities"],
         )
 
         return raw_predictions

tiatoolbox/models/models_abc.py

@@ -102,15 +102,17 @@ def forward(
 
     @staticmethod
     @abstractmethod
-    def infer_batch(model: nn.Module, batch_data: np.ndarray, *, device: str) -> dict:
+    def infer_batch(
+        model: nn.Module, batch_data: np.ndarray | torch.Tensor, *, device: str
+    ) -> dict:
         """Run inference on an input batch.
 
         Contains logic for forward operation as well as I/O aggregation.
 
         Args:
             model (nn.Module):
                 PyTorch defined model.
-            batch_data (np.ndarray):
+            batch_data (np.ndarray | torch.Tensor):
                 A batch of data generated by
                 `torch.utils.data.DataLoader`.
             device (str):