MIL component to extract patches

myron · myron · commit 463b1593186a · 2021-11-11T01:21:05.000-08:00
Signed-off-by: myron &lt;amyronenko@nvidia.com&gt;
diff --git a/monai/apps/pathology/transforms/__init__.py b/monai/apps/pathology/transforms/__init__.py
@@ -9,8 +9,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .spatial.array import SplitOnGrid
-from .spatial.dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict
+from .spatial.array import SplitOnGrid, TileOnGrid
+from .spatial.dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict, TileOnGridd, TileOnGridD, TileOnGridDict
 from .stain.array import ExtractHEStains, NormalizeHEStains
 from .stain.dictionary import (
     ExtractHEStainsd,
diff --git a/monai/apps/pathology/transforms/spatial/__init__.py b/monai/apps/pathology/transforms/spatial/__init__.py
@@ -9,5 +9,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .array import SplitOnGrid
-from .dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict
+from .array import SplitOnGrid, TileOnGrid
+from .dictionary import SplitOnGridd, SplitOnGridD, SplitOnGridDict, TileOnGridd, TileOnGridD, TileOnGridDict
diff --git a/monai/apps/pathology/transforms/spatial/array.py b/monai/apps/pathology/transforms/spatial/array.py
@@ -9,13 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Tuple, Union
+from typing import Optional, Sequence, Tuple, Union
 
+import numpy as np
 import torch
+from numpy.lib.stride_tricks import as_strided
 
-from monai.transforms.transform import Transform
+from monai.transforms.transform import Randomizable, Transform
 
-__all__ = ["SplitOnGrid"]
+__all__ = ["SplitOnGrid", "TileOnGrid"]
 
 
 class SplitOnGrid(Transform):
@@ -73,3 +75,160 @@ def get_params(self, image_size):
         )
 
         return patch_size, steps
+
+
+class TileOnGrid(Randomizable, Transform):
+    """
+    Tile the 2D image into patches on a grid and maintain a subset of it.
+    This transform works only with np.ndarray inputs for 2D images.
+
+    Args:
+        tile_count: number of tiles to extract
+        tile_size: size of the square tile
+            Defaults to ``256``.
+        tile_step: step size
+            Defaults to None (same as tile_size)
+        tile_all: Extract all non-background tiles, instead of tile_count.
+            Defaults to ``False``.
+        tile_random_offset: Randomize position of tile grid, instead of starting from the top-left corner
+            Defaults to ``False``.
+        tile_pad_full: pad image to the size evenly divisible by tile_size
+            Defaults to ``False``.
+        tile_background_val: the background constant (e.g. 255 for white background)
+            Defaults to ``255``.
+        tile_filter_mode: mode must be in ["min", "max", None]. If total number of tiles is more then tile_size,
+            then sort by intensity sum, and take the smallest (for min), largest (for max) or random (for None) subset
+            Defaults to ``min`` (which assumes background is white, high value)
+
+    """
+
+    def __init__(
+        self,
+        tile_count: Optional[int] = None,
+        tile_size: int = 256,
+        tile_step: Optional[int] = None,
+        tile_all: bool = False,
+        tile_random_offset: bool = False,
+        tile_pad_full: bool = False,
+        tile_background_val: int = 255,
+        tile_filter_mode: Optional[str] = "min",
+    ):
+        self.tile_count = tile_count
+        self.tile_size = tile_size
+        self.tile_step = tile_step
+        self.tile_all = tile_all
+        self.tile_random_offset = tile_random_offset
+        self.tile_pad_full = tile_pad_full
+        self.tile_background_val = tile_background_val
+        self.tile_filter_mode = tile_filter_mode
+
+        if self.tile_count is None:
+            self.tile_count = max((44 * 256 ** 2) // (tile_size ** 2), 1)
+
+        if self.tile_step is None:
+            self.tile_step = self.tile_size  # non-overlapping grid
+
+        self.random_offset = (0, 0)
+        self.random_idxs = [0]
+
+    def randomize(self, img_size: Sequence[int]) -> None:
+
+        c, h, w = img_size
+        tile_count: int = self.tile_count  # type: ignore
+        tile_step: int = self.tile_step  # type: ignore
+
+        if self.tile_random_offset:
+            pad_h = h % self.tile_size
+            pad_w = w % self.tile_size
+            if pad_h > 0 and pad_w > 0:
+                self.random_offset = (self.R.randint(pad_h), self.R.randint(pad_w))
+                h = h - self.random_offset[0]
+                w = w - self.random_offset[1]
+            else:
+                self.random_offset = (0, 0)
+
+        if self.tile_pad_full:
+            pad_h = (self.tile_size - h % self.tile_size) % self.tile_size
+            pad_w = (self.tile_size - w % self.tile_size) % self.tile_size
+            h = h + pad_h
+            w = w + pad_w
+
+        h_n = (h - self.tile_size + tile_step) // tile_step
+        w_n = (w - self.tile_size + tile_step) // tile_step
+        tile_total = h_n * w_n
+
+        if tile_total > tile_count:
+            self.random_idxs = self.R.choice(range(tile_total), tile_count, replace=False)  # type: ignore
+        else:
+            self.random_idxs = [0]  # type: ignore
+
+    def __call__(self, image: np.ndarray) -> np.ndarray:
+
+        # add random offset
+        self.randomize(img_size=image.shape)
+        tile_count: int = self.tile_count  # type: ignore
+        tile_step: int = self.tile_step  # type: ignore
+
+        if self.tile_random_offset and self.random_offset is not None:
+            image = image[:, self.random_offset[0] :, self.random_offset[1] :]
+
+        # pad to full size, divisible by tile_size
+        if self.tile_pad_full:
+            c, h, w = image.shape
+            pad_h = (self.tile_size - h % self.tile_size) % self.tile_size
+            pad_w = (self.tile_size - w % self.tile_size) % self.tile_size
+            image = np.pad(
+                image,
+                [[0, 0], [pad_h // 2, pad_h - pad_h // 2], [pad_w // 2, pad_w - pad_w // 2]],
+                constant_values=self.tile_background_val,
+            )
+
+        # extact tiles (new way)
+        xstep, ystep = tile_step, tile_step
+        xsize, ysize = self.tile_size, self.tile_size
+        clen, xlen, ylen = image.shape
+        cstride, xstride, ystride = image.strides
+        llw = as_strided(
+            image,
+            shape=((xlen - xsize) // xstep + 1, (ylen - ysize) // ystep + 1, clen, xsize, ysize),
+            strides=(xstride * xstep, ystride * ystep, cstride, xstride, ystride),
+            writeable=False,
+        )
+        image = llw.reshape(-1, clen, xsize, ysize)
+
+        # if keep all patches
+        if self.tile_all:
+            # retain only patches with significant foreground content to speed up inference
+            # FYI, this returns a variable number of tiles, so the batch_size much be 1 (per gpu). Used during inference
+            thresh = 0.999 * 3 * self.tile_background_val * self.tile_size * self.tile_size
+            if self.tile_filter_mode == "min":
+                # default, keep non-background tiles (small values)
+                idxs = np.argwhere(image.sum(axis=(1, 2, 3)) < thresh)
+                image = image[idxs.reshape(-1)]
+            elif self.tile_filter_mode == "max":
+                idxs = np.argwhere(image.sum(axis=(1, 2, 3)) >= thresh)
+                image = image[idxs.reshape(-1)]
+
+        else:
+            if len(image) >= tile_count:
+
+                if self.tile_filter_mode == "min":
+                    # default, keep non-background tiles (smallest values)
+                    idxs = np.argsort(image.sum(axis=(1, 2, 3)))[:tile_count]
+                    image = image[idxs]
+                elif self.tile_filter_mode == "max":
+                    idxs = np.argsort(image.sum(axis=(1, 2, 3)))[-tile_count:]
+                    image = image[idxs]
+                elif len(image) > tile_count:
+                    # random subset (more appropriate for WSIs without distinct background)
+                    if self.random_idxs is not None:
+                        image = image[self.random_idxs]
+
+            else:
+                image = np.pad(
+                    image,
+                    [[0, tile_count - len(image)], [0, 0], [0, 0], [0, 0]],
+                    constant_values=self.tile_background_val,
+                )
+
+        return image
diff --git a/monai/apps/pathology/transforms/spatial/dictionary.py b/monai/apps/pathology/transforms/spatial/dictionary.py
@@ -9,16 +9,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, Hashable, Mapping, Optional, Tuple, Union
+import copy
+from typing import Any, Dict, Hashable, List, Mapping, Optional, Tuple, Union
 
+import numpy as np
 import torch
 
 from monai.config import KeysCollection
-from monai.transforms.transform import MapTransform
+from monai.transforms.transform import MapTransform, Randomizable
 
-from .array import SplitOnGrid
+from .array import SplitOnGrid, TileOnGrid
 
-__all__ = ["SplitOnGridd", "SplitOnGridD", "SplitOnGridDict"]
+__all__ = ["SplitOnGridd", "SplitOnGridD", "SplitOnGridDict", "TileOnGridd", "TileOnGridD", "TileOnGridDict"]
 
 
 class SplitOnGridd(MapTransform):
@@ -53,4 +55,83 @@ def __call__(self, data: Mapping[Hashable, torch.Tensor]) -> Dict[Hashable, torc
         return d
 
 
+class TileOnGridd(Randomizable, MapTransform):
+    """
+    Tile the 2D image into patches on a grid and maintain a subset of it.
+    This transform works only with np.ndarray inputs for 2D images.
+
+    Args:
+        tile_count: number of tiles to extract
+        tile_size: size of the square tile
+            Defaults to ``256``.
+        tile_step: step size
+            Defaults to None (same as tile_size)
+        tile_all: Extract all non-background tiles, instead of tile_count.
+            Defaults to ``False``.
+        tile_random_offset: Randomize position of tile grid, instead of starting from the top-left corner
+            Defaults to ``False``.
+        tile_pad_full: pad image to the size evenly divisible by tile_size
+            Defaults to ``False``.
+        tile_background_val: the background constant (e.g. 255 for white background)
+            Defaults to ``255``.
+        tile_filter_mode: mode must be in ["min", "max", None]. If total number of tiles is more then tile_size,
+            then sort by intensity sum, and take the smallest (for min), largest (for max) or random (for None) subset
+            Defaults to ``min`` (which assumes background is white, high value)
+        return_list_of_dicts: return each tile in a separate dictionary, as a list of dicts
+            Defaults to ``False``
+
+    """
+
+    def __init__(
+        self,
+        keys: KeysCollection,
+        tile_count: Optional[int] = None,
+        tile_size: int = 256,
+        tile_step: Optional[int] = None,
+        tile_all: bool = False,
+        tile_random_offset: bool = False,
+        tile_pad_full: bool = False,
+        tile_background_val: int = 255,
+        tile_filter_mode: str = "min",
+        allow_missing_keys: bool = False,
+        return_list_of_dicts: bool = False,
+    ):
+        super().__init__(keys, allow_missing_keys)
+
+        self.return_list_of_dicts = return_list_of_dicts
+        self.seed = None
+
+        self.splitter = TileOnGrid(
+            tile_count=tile_count,
+            tile_size=tile_size,
+            tile_step=tile_step,
+            tile_all=tile_all,
+            tile_random_offset=tile_random_offset,
+            tile_pad_full=tile_pad_full,
+            tile_background_val=tile_background_val,
+            tile_filter_mode=tile_filter_mode,
+        )
+
+    def randomize(self, data: Any = None) -> None:
+        self.seed = self.R.randint(10000)  # type: ignore
+
+    def __call__(self, data: Mapping[Hashable, np.ndarray]) -> Union[Dict[Hashable, np.ndarray], List[Dict]]:
+
+        self.randomize()
+
+        d = dict(data)
+        for key in self.key_iterator(d):
+            self.splitter.set_random_state(seed=self.seed)  # same random seed for all keys
+            d[key] = self.splitter(d[key])
+
+        if self.return_list_of_dicts:
+            d_list = []
+            for i in range(len(d[self.keys[0]])):
+                d_list.append({k: d[k][i] if k in self.keys else copy.deepcopy(d[k]) for k in d.keys()})
+            d = d_list  # type: ignore
+
+        return d
+
+
 SplitOnGridDict = SplitOnGridD = SplitOnGridd
+TileOnGridDict = TileOnGridD = TileOnGridd
diff --git a/tests/test_tile_on_grid.py b/tests/test_tile_on_grid.py
diff --git a/tests/test_tile_on_grid_dict.py b/tests/test_tile_on_grid_dict.py
