feat(datamodules): add h5 writing to pannuke

Author: okunator

cellseg_models_pytorch/datamodules/lizard_datamodule.py

@@ -6,8 +6,8 @@
 
 try:
     from ..datasets import SegmentationFolderDataset, SegmentationHDF5Dataset
-    from ..datasets.dataset_writers.folder_writer import SlidingWindowFolderWriter
-    from ..datasets.dataset_writers.hdf5_writer import SlidingWindowHDF5Writer
+    from ..datasets.dataset_writers.folder_writer import FolderWriter
+    from ..datasets.dataset_writers.hdf5_writer import HDF5Writer
     from ._basemodule import BaseDataModule
     from .downloader import SimpleDownloader
 except ModuleNotFoundError:
@@ -50,7 +50,7 @@ def __init__(
         Parameters
         ----------
             save_dir : str
-                Path to directory where the pannuke data will be saved.
+                Path to directory where the lizard data will be saved.
             fold_split : Dict[str, int]
                 Defines how the folds are split into train, valid, and test sets.
                 E.g. {"train": 1, "valid": 2, "test": 3}
@@ -82,18 +82,20 @@ def __init__(
             >>> from cellseg_models_pytorch.datamodules import LizardDataModule
 
             >>> fold_split = {"train": 1, "valid": 2, "test": 3}
-            >>> save_dir = Path.home() / "pannuke"
+            >>> save_dir = Path.home() / "lizard"
             >>> lizard_module = LizardDataModule(
                     save_dir=save_dir,
                     fold_split=fold_split,
                     inst_transforms=["dist", "stardist"],
                     img_transforms=["blur", "hue_sat"],
                     normalization="percentile",
-                    dataset_type="hdf5"
+                    dataset_type="hdf5",
+                    patch_size=(320, 320),
+                    stride=128
                 )
 
             >>> # lizard_module.download(save_dir) # just the downloading
-            >>> lizard_module.prepare_data(do_patching=True) # downloading & processing
+            >>> lizard_module.prepare_data(tiling=True) # downloading & processing
         """
         super().__init__(batch_size, num_workers)
         self.save_dir = Path(save_dir)
@@ -115,7 +117,7 @@ def __init__(
 
     @property
     def type_classes(self) -> Dict[str, int]:
-        """Pannuke cell type classes."""
+        """Lizard cell type classes."""
         return {
             "bg": 0,
             "neutrophil": 1,
@@ -138,7 +140,7 @@ def download(root: str) -> None:
         SimpleDownloader.download(url, root)
         LizardDataModule.extract_zips(root, rm=True)
 
-    def prepare_data(self, rm_orig: bool = False, do_patching: bool = True) -> None:
+    def prepare_data(self, rm_orig: bool = False, tiling: bool = True) -> None:
         """Prepare the lizard datasets.
 
         1. Download lizard folds from:
@@ -151,9 +153,9 @@ def prepare_data(self, rm_orig: bool = False, do_patching: bool = True) -> None:
             rm_orig : bool, default=False
                 After processing all the files, If True, removes the original
                 un-processed files.
-            do_patching : bool, default=True
-                Flag, whether to do patching at all. Can be used if you only want to
-                download and split the data and then work it out on your own.
+            tiling : bool, default=True
+                Flag, whether to cut images into tiles. Can be set to False if you only
+                want to download and split the data and then work it out on your own.
         """
         folders_found = [
             d.name
@@ -222,7 +224,7 @@ def prepare_data(self, rm_orig: bool = False, do_patching: bool = True) -> None:
                 if "lizard" in d.name.lower() or "macosx" in d.name.lower():
                     shutil.rmtree(d)
 
-        if do_patching and not patches_found:
+        if tiling and not patches_found:
             print("Patch the data... This will take a while...")
             for phase in self.fold_split.keys():
                 save_im_dir = self.save_dir / phase / "images"
@@ -231,7 +233,7 @@ def prepare_data(self, rm_orig: bool = False, do_patching: bool = True) -> None:
                 if self.dataset_type == "hdf5":
                     sdir = self.save_dir / phase / f"{phase}_patches"
                     sdir.mkdir(parents=True, exist_ok=True)
-                    writer = SlidingWindowHDF5Writer(
+                    writer = HDF5Writer(
                         in_dir_im=save_im_dir,
                         in_dir_mask=save_mask_dir,
                         save_dir=sdir,
@@ -245,7 +247,7 @@ def prepare_data(self, rm_orig: bool = False, do_patching: bool = True) -> None:
                     sdir_mask = self.save_dir / phase / f"{phase}_mask_patches"
                     sdir_im.mkdir(parents=True, exist_ok=True)
                     sdir_mask.mkdir(parents=True, exist_ok=True)
-                    writer = SlidingWindowFolderWriter(
+                    writer = FolderWriter(
                         in_dir_im=save_im_dir,
                         in_dir_mask=save_mask_dir,
                         save_dir_im=sdir_im,
@@ -254,7 +256,7 @@ def prepare_data(self, rm_orig: bool = False, do_patching: bool = True) -> None:
                         stride=self.stride,
                         transforms=["rigid"],
                     )
-                writer.write(pre_proc=self._process_label, msg=phase)
+                writer.write(tiling=True, pre_proc=self._process_label, msg=phase)
         else:
             print(
                 "Found processed Lizard data. "

cellseg_models_pytorch/datamodules/pannuke_datamodule.py

@@ -8,7 +8,8 @@
 from ..utils import FileHandler, fix_duplicates
 
 try:
-    from ..datasets import SegmentationFolderDataset
+    from ..datasets import SegmentationFolderDataset, SegmentationHDF5Dataset
+    from ..datasets.dataset_writers.hdf5_writer import HDF5Writer
     from ._basemodule import BaseDataModule
     from .downloader import SimpleDownloader
 except ModuleNotFoundError:
@@ -26,6 +27,7 @@ def __init__(
         fold_split: Dict[str, int],
         img_transforms: List[str],
         inst_transforms: List[str],
+        dataset_type: str = "folder",
         normalization: str = None,
         batch_size: int = 8,
         num_workers: int = 8,
@@ -65,6 +67,8 @@ def __init__(
                 A list containg all the transformations that are applied to only the
                 instance labelled masks. Allowed ones: "cellpose", "contour", "dist",
                 "edgeweight", "hovernet", "omnipose", "smooth_dist", "binarize"
+            dataset_type : str, default="folder"
+                The dataset type. One of "folder", "hdf5".
             normalization : str, optional
                 Apply img normalization after all the transformations. One of "minmax",
                 "norm", "percentile", None.
@@ -107,6 +111,14 @@ def __init__(
         self.normalization = normalization
         self.kwargs = kwargs if kwargs is not None else {}
 
+        if dataset_type not in ("folder", "hdf5"):
+            raise ValueError(
+                f"Illegal `dataset_type` arg. Got {dataset_type}. "
+                f"Allowed: {('folder', 'hdf5')}"
+            )
+
+        self.dataset_type = dataset_type
+
     @property
     def type_classes(self) -> Dict[str, int]:
         """Pannuke cell type classes."""
@@ -127,7 +139,7 @@ def download(root: str) -> None:
             SimpleDownloader.download(url, root)
         PannukeDataModule.extract_zips(root, rm=True)
 
-    def prepare_data(self, rm_orig: bool = True) -> None:
+    def prepare_data(self, rm_orig: bool = False) -> None:
         """Prepare the pannuke datasets.
 
         1. Download pannuke folds from:
@@ -167,6 +179,18 @@ def prepare_data(self, rm_orig: bool = True) -> None:
                 self._process_pannuke_fold(
                     fold_paths, save_im_dir, save_mask_dir, fold_ix, phase
                 )
+
+                if self.dataset_type == "hdf5":
+                    writer = HDF5Writer(
+                        in_dir_im=save_im_dir,
+                        in_dir_mask=save_mask_dir,
+                        save_dir=self.save_dir / phase,
+                        file_name=f"pannuke_{phase}.h5",
+                        patch_size=None,
+                        stride=None,
+                        transforms=None,
+                    )
+                    writer.write(tiling=False, msg=phase)
         else:
             print(
                 "Found processed pannuke data. "
@@ -178,31 +202,45 @@ def prepare_data(self, rm_orig: bool = True) -> None:
                 if "fold" in d.name.lower():
                     shutil.rmtree(d)
 
+    def _get_path(self, phase: str, dstype: str, is_mask: bool = False) -> Path:
+        if dstype == "hdf5":
+            p = self.save_dir / phase / f"pannuke_{phase}.h5"
+        else:
+            dtype = "labels" if is_mask else "images"
+            p = self.save_dir / phase / dtype
+
+        return p
+
     def setup(self, stage: Optional[str] = None) -> None:
         """Set up the train, valid, and test datasets."""
-        self.trainset = SegmentationFolderDataset(
-            path=self.save_dir / "train" / "images",
-            mask_path=self.save_dir / "train" / "labels",
+        if self.dataset_type == "hdf5":
+            DS = SegmentationHDF5Dataset
+        else:
+            DS = SegmentationFolderDataset
+
+        self.trainset = DS(
+            path=self._get_path("train", self.dataset_type, is_mask=False),
+            mask_path=self._get_path("train", self.dataset_type, is_mask=True),
             img_transforms=self.img_transforms,
             inst_transforms=self.inst_transforms,
             return_sem=False,
             normalization=self.normalization,
             **self.kwargs,
         )
 
-        self.validset = SegmentationFolderDataset(
-            path=self.save_dir / "valid" / "images",
-            mask_path=self.save_dir / "valid" / "labels",
+        self.validset = DS(
+            path=self._get_path("valid", self.dataset_type, is_mask=False),
+            mask_path=self._get_path("valid", self.dataset_type, is_mask=True),
             img_transforms=self.img_transforms,
             inst_transforms=self.inst_transforms,
             return_sem=False,
             normalization=self.normalization,
             **self.kwargs,
         )
 
-        self.testset = SegmentationFolderDataset(
-            path=self.save_dir / "test" / "images",
-            mask_path=self.save_dir / "test" / "labels",
+        self.testset = DS(
+            path=self._get_path("test", self.dataset_type, is_mask=False),
+            mask_path=self._get_path("test", self.dataset_type, is_mask=True),
             img_transforms=self.img_transforms,
             inst_transforms=self.inst_transforms,
             return_sem=False,
@@ -256,7 +294,7 @@ def _process_pannuke_fold(
                     inst_map = self._get_inst_map(temp_mask[..., 0:5])
 
                     fn_mask = Path(save_mask_dir / name).with_suffix(".mat")
-                    FileHandler.write_mask(fn_mask, inst_map, type_map)
+                    FileHandler.write_mat(fn_mask, inst_map, type_map)
                     pbar.update(1)
 
     def _get_type_map(self, pannuke_mask: np.ndarray) -> np.ndarray:

changelog.d/20230113_171853_oskari.lehtonen.md

@@ -0,0 +1,3 @@
+## Features
+
+- Add option to write pannuke dataset to h5 db in `PannukeDataModule` and `LizardDataModule`.