fix(datasets): fix dataset writers shortcomings

Author: okunator

cellseg_models_pytorch/datasets/dataset_writers/__init__.py

@@ -1,4 +1,4 @@
 from ._base_writer import BaseWriter
-from .folder_writer import SlidingWindowFolderWriter
+from .folder_writer import FolderWriter
 
-__all__ = ["BaseWriter", "SlidingWindowFolderWriter"]
+__all__ = ["BaseWriter", "FolderWriter"]

cellseg_models_pytorch/datasets/dataset_writers/_base_writer.py

@@ -6,7 +6,7 @@
 from pathos.multiprocessing import ThreadPool as Pool
 
 from ...transforms.albu_transforms import IMG_TRANSFORMS, compose
-from ...utils import FileHandler, TilerStitcher, fix_duplicates
+from ...utils import FileHandler, TilerStitcher, fix_duplicates, remap_label
 
 __all__ = ["BaseWriter"]
 
@@ -18,48 +18,57 @@ class BaseWriter(ABC):
     def __init__(
         self,
         in_dir_im: str,
-        in_dir_mask: str,
-        patch_size: Tuple[int, int],
-        stride: int,
+        in_dir_mask: str = None,
+        patch_size: Tuple[int, int] = None,
+        stride: int = None,
         transforms: Optional[List[str]] = None,
         **kwargs,
     ) -> None:
         """Init base class for sliding window data writers."""
-        self.im_dir = Path(in_dir_im)
-        self.mask_dir = Path(in_dir_mask)
         self.stride = stride
+
+        if isinstance(patch_size, int):
+            patch_size = (patch_size, patch_size)
+
         self.patch_size = patch_size
+        self.im_dir = Path(in_dir_im)
 
+        # Imgs
         if not self.im_dir.exists():
             raise ValueError(f"folder: {self.im_dir} does not exist")
 
         if not self.im_dir.is_dir():
             raise ValueError(f"path: {self.im_dir} is not a folder")
 
-        if not all([f.suffix in IMG_SUFFIXES for f in self.im_dir.iterdir()]):
-            raise ValueError(
-                f"files formats in given folder need to be in {IMG_SUFFIXES}"
-            )
+        im_files = []
+        for types in IMG_SUFFIXES:
+            im_files.extend(self.im_dir.glob(f"*{types}"))
+        self.fnames_im = sorted(im_files)
 
-        if not self.mask_dir.exists():
-            raise ValueError(f"folder: {self.mask_dir} does not exist")
+        # Masks
+        self.mask_dir = in_dir_mask
+        self.fnames_mask = None
+        if in_dir_mask is not None:
+            self.mask_dir = Path(in_dir_mask)
 
-        if not self.mask_dir.is_dir():
-            raise ValueError(f"path: {self.mask_dir} is not a folder")
+            if not self.mask_dir.exists():
+                raise ValueError(f"folder: {self.mask_dir} does not exist")
 
-        if not all([f.suffix in MASK_SUFFIXES for f in self.mask_dir.iterdir()]):
-            raise ValueError(
-                f"files formats in given folder need to be in {MASK_SUFFIXES}"
-            )
+            if not self.mask_dir.is_dir():
+                raise ValueError(f"path: {self.mask_dir} is not a folder")
 
-        self.fnames_im = sorted(self.im_dir.glob("*"))
-        self.fnames_mask = sorted(self.mask_dir.glob("*"))
-        if len(self.fnames_im) != len(self.fnames_mask):
-            raise ValueError(
-                f"Found different number of files in {self.im_dir.as_posix()} and "
-                f"{self.mask_dir.as_posix()}."
-            )
+            mask_files = []
+            for types in MASK_SUFFIXES:
+                mask_files.extend(self.mask_dir.glob(f"*{types}"))
+            self.fnames_mask = sorted(mask_files)
 
+            if len(self.fnames_im) != len(self.fnames_mask):
+                raise ValueError(
+                    f"Found different number of files in {self.im_dir.as_posix()} and "
+                    f"{self.mask_dir.as_posix()}."
+                )
+
+        # Transformations
         self.transforms = None
         if transforms is not None:
             allowed = list(IMG_TRANSFORMS.keys())
@@ -77,78 +86,200 @@ def write(self):
         """Patch images and mask to and write them to disk."""
         raise NotImplementedError
 
-    def _get_tiles(
+    def get_array(
         self,
         img_path: Union[str, Path],
-        mask_path: Union[str, Path],
+        mask_path: Optional[Union[str, Path]] = None,
+        tiling: Optional[bool] = False,
+        pre_proc: Optional[Callable] = None,
+    ) -> Tuple[np.ndarray, Union[None, Dict[str, np.ndarray]]]:
+        """Pipeline that (optionally) patches and transforms input images and masks.
+
+        Parameters
+        ----------
+            img_path : str or Path
+                Path to an image file.
+            mask_path : str or Path, optional
+                Path to a .mat mask file.
+            tiling : bool, default=False, optional
+                Flag, whether to do tiling on the images (and masks).
+            pre_proc : Callable, optional
+                A pre-processing function that can be used to pre-process given input
+                masks before the pipeline.
+
+        Raises
+        ------
+            ValueError if self.stride or self.patch_size are not set to integer values.
+
+        Returns
+        -------
+            Tuple[np.ndarray, Union[None, Dict[str, np.ndarray]]]:
+                The processed image & masks. If `mask_path=None`, returns no masks.
+                Img shape w/o tiling: (H, W, C). Dtype: uint8.
+                Img shape w/ tiling: (N, pH, pW, C). Dtype: uint8.
+                Masks w/o tiling: Shapes: (H, W). Dtypes: int32.
+                Masks w/ tiling: Shapes: (N, pH, pW). Dtypes: int32.
+        """
+        im, masks = self._read_files(img_path, mask_path, pre_proc)
+
+        # do tiling first if flag set to True
+        if tiling:
+            if not isinstance(self.stride, int) and not isinstance(
+                self.patch_size, int
+            ):
+                raise ValueError(
+                    "`self.stride` and `self.patch_size` need to be integers. Got: "
+                    f"self.stride={self.stride}, self.patch_size={self.patch_size}"
+                )
+
+            im, masks = self._get_tiles(im, masks)
+
+            if masks is not None:
+                if "inst_map" in masks.keys():
+                    masks["inst_map"] = self._fix_instances_tiles(masks["inst_map"])
+
+            if self.transforms is not None:
+                im, masks = self._transform_tiles(im, masks)
+        else:
+            if masks is not None:
+                if "inst_map" in masks.keys():
+                    masks["inst_map"] = self._fix_instances_one(masks["inst_map"])
+
+            if self.transforms is not None:
+                im, masks = self._transform_one(im, masks)
+
+        return im, masks
+
+    def _fix_instances_one(self, inst_map: np.ndarray) -> np.ndarray:
+        """Fix duplicate instances and remap instance labels."""
+        return remap_label(fix_duplicates(inst_map))
+
+    def _read_files(
+        self,
+        img_path: Union[str, Path],
+        mask_path: Union[str, Path] = None,
         pre_proc: Callable = None,
-    ) -> Dict[str, np.ndarray]:
-        """Read one image and corresponding masks and do tiling on them."""
-        # im, masks = self._get_arrays()
+    ) -> Tuple[np.ndarray, Union[None, Dict[str, np.ndarray]]]:
+        """Read image and corresponding masks if there are such."""
         im = FileHandler.read_img(img_path)
-        masks = FileHandler.read_mat(mask_path, return_all=True)
 
-        if pre_proc is not None:
-            masks = pre_proc(masks)
+        masks = None
+        if mask_path is not None:
+            masks = FileHandler.read_mat(mask_path, return_all=True)
+
+            if pre_proc is not None:
+                masks = pre_proc(masks)
 
-        inst = None
-        types = None
-        sem = None
-        if "inst_map" in masks.keys():
-            inst = masks["inst_map"]
-        if "type_map" in masks.keys():
-            types = masks["type_map"]
-        if "sem_map" in masks.keys():
-            sem = masks["sem_map"]
+            masks = {
+                key: arr
+                for key, arr in masks.items()
+                if key in ("inst_map", "type_map", "sem_map")
+            }
 
+        return im, masks
+
+    def _get_tiles(
+        self,
+        im: np.ndarray,
+        masks: Union[Dict[str, np.ndarray], None] = None,
+    ) -> Tuple[Dict[str, np.ndarray], Union[Dict[str, np.ndarray], None]]:
+        """Do tiling on an image and corresponding masks if there are such."""
+        # Init Tilers
         im_tiler = TilerStitcher(
             im_shape=im.shape, patch_shape=self.patch_size + (3,), stride=self.stride
         )
+        im_tiles = im_tiler.patch(im)
+
+        # Tile masks if there are such.
+        mask_tiles = None
+        if masks is not None:
+            mask_tiles = {}
+            inst = None
+            types = None
+            sem = None
+            if "inst_map" in masks.keys():
+                inst = masks["inst_map"]
+            if "type_map" in masks.keys():
+                types = masks["type_map"]
+            if "sem_map" in masks.keys():
+                sem = masks["sem_map"]
+
+            mask_tiler = TilerStitcher(
+                im_shape=inst.shape,
+                patch_shape=self.patch_size + (1,),
+                stride=self.stride,
+            )
 
-        mask_tiler = TilerStitcher(
-            im_shape=inst.shape, patch_shape=self.patch_size + (1,), stride=self.stride
-        )
+            if inst is not None:
+                mask_tiles["inst_map"] = mask_tiler.patch(inst).squeeze()
+            if types is not None:
+                mask_tiles["type_map"] = mask_tiler.patch(types).squeeze()
+            if sem is not None:
+                mask_tiles["sem_map"] = mask_tiler.patch(sem).squeeze()
+
+        return im_tiles, mask_tiles
 
-        tiles = {}
-        tiles["image"] = im_tiler.patch(im)
-        if inst is not None:
-            tiles["inst_map"] = self._fix_duplicates(mask_tiler.patch(inst).squeeze())
-        if types is not None:
-            tiles["type_map"] = mask_tiler.patch(types).squeeze()
-        if sem is not None:
-            tiles["sem_map"] = mask_tiler.patch(sem).squeeze()
+    def _transform_one(
+        self, im: np.ndarray, masks: Dict[str, np.ndarray] = None
+    ) -> Tuple[np.ndarray, Union[Dict[str, np.ndarray], None]]:
+        """Transform an image and corresponding mask if there is one."""
+        if masks is not None:
+            mask_names = [name for name in masks.keys()]
+            masks = [mask for mask in masks.values()]
+            out = self.transforms(image=im, masks=masks)
+            masks = {n: mask for n, mask in zip(mask_names, out["masks"])}
+        else:
+            out = self.transforms(image=im)
 
-        if self.transforms is not None:
-            tiles = self._transform(tiles)
+        im = out["image"]
 
-        return tiles
+        return im, masks
 
-    def _transform(self, tiles: Dict[str, np.ndarray]) -> np.ndarray:
+    def _transform_tiles(
+        self,
+        im_tiles: np.ndarray,
+        mask_tiles: Union[Dict[str, np.ndarray], None] = None,
+    ) -> Tuple[np.ndarray, Union[Dict[str, np.ndarray], None]]:
         """Apply transformations to the tiles one by one."""
-        n_tiles = tiles["image"].shape[0]
-        masks = [arr for key, arr in tiles.items() if key != "image"]
-        mask_names = [key for key in tiles.keys() if key != "image"]
+        n_tiles = im_tiles.shape[0]
+        out_im_tiles = []
+
+        out_mask_tiles = None
+        if mask_tiles is not None:
+            mask_names = [key for key in mask_tiles.keys()]
+            out_mask_tiles = {k: [] for k in mask_tiles.keys()}
 
-        out_tiles = {k: [] for k in tiles.keys()}
         for i in range(n_tiles):
-            m = [mask[i] for mask in masks]
-            out = self.transforms(image=tiles["image"][i], masks=m)
-            out_tiles["image"].append(out["image"])
+            # Get one img tile
+            im = im_tiles[i]
+
+            # Get one set of mask tiles
+            masks = None
+            if mask_tiles is not None:
+                masks = {n: mask_tiles[n][i] for n in mask_names}
+
+            # transform imgs & masks
+            im_tr, masks_tr = self._transform_one(im, masks)
 
-            for j, mname in enumerate(mask_names):
-                out_tiles[mname].append(out["masks"][j])
+            out_im_tiles.append(im_tr)
+            if mask_tiles is not None:
+                for mask_name in mask_names:
+                    out_mask_tiles[mask_name].append(masks_tr[mask_name])
 
-        for k, mask in out_tiles.items():
-            out_tiles[k] = np.array(mask)
+        # convert list of 2D-arrays to np.ndarray
+        out_im_tiles = np.array(out_im_tiles)
+        if mask_tiles is not None:
+            for k, arr in out_mask_tiles.items():
+                out_mask_tiles[k] = np.array(arr)
 
-        return out_tiles
+        return out_im_tiles, out_mask_tiles
 
-    def _fix_duplicates(self, patches_inst: np.ndarray) -> np.ndarray:
-        """Fix repeatded labels in a patched instance labelled mask."""
+    def _fix_instances_tiles(self, patches_inst: np.ndarray) -> np.ndarray:
+        """Fix repeated labels and remap them in a patched instance labelled mask."""
         insts = []
 
         for i in range(patches_inst.shape[0]):
-            insts.append(fix_duplicates(patches_inst[i]))
+            insts.append(self._fix_instances_one(patches_inst[i]))
 
         insts = np.array(insts)