fix(datasets): move infer datset, add err catch

Author: okunator

cellseg_models_pytorch/datasets/__init__.py

@@ -1,4 +1,4 @@
-from .folder_dataset import FolderDataset
 from .folder_dataset_train import SegmentationFolderDataset
+from .hdf5_dataset import SegmentationHDF5Dataset
 
-__all__ = ["FolderDataset", "SegmentationFolderDataset"]
+__all__ = ["FolderDataset", "SegmentationFolderDataset", "SegmentationHDF5Dataset"]

cellseg_models_pytorch/datasets/_base_dataset.py

@@ -3,14 +3,20 @@
 import numpy as np
 from torch.utils.data import Dataset
 
-from ..transforms import (
-    IMG_TRANSFORMS,
-    INST_TRANSFORMS,
-    NORM_TRANSFORMS,
-    apply_each,
-    compose,
-    to_tensorv3,
-)
+try:
+    from ..transforms.albu_transforms import (
+        IMG_TRANSFORMS,
+        INST_TRANSFORMS,
+        NORM_TRANSFORMS,
+        apply_each,
+        compose,
+        to_tensorv3,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "To use the `csmp.dataset` module, the albumentations lib is needed. "
+        "Install with `pip install albumentations`"
+    )
 
 __all__ = ["TrainDatasetBase"]
 

cellseg_models_pytorch/datasets/tests/test_dataset.py

@@ -1,47 +1,47 @@
-# import pytest
-# import torch
-
-# from cellseg_models_pytorch.datasets import SegmentationHDF5Dataset
-
-# img_transforms = ["rigid", "blur"]
-# inst_transforms = ["smooth_dist"]
-
-
-# @pytest.mark.optional
-# @pytest.mark.parametrize("return_inst", [True, False])
-# @pytest.mark.parametrize("return_type", [True, False])
-# @pytest.mark.parametrize("return_sem", [True, False])
-# @pytest.mark.parametrize("normalization", [None, "minmax"])
-# def test_hdf5_dataset(hdf5db, return_inst, return_type, return_sem, normalization):
-#     ds = SegmentationHDF5Dataset(
-#         path=hdf5db,
-#         img_transforms=img_transforms,
-#         inst_transforms=inst_transforms,
-#         normalization=normalization,
-#         return_inst=return_inst,
-#         return_type=return_type,
-#         return_sem=return_sem,
-#     )
-
-#     out = next(iter(ds))
-
-#     if return_inst:
-#         assert "inst" in out.keys()
-#         assert out["inst"].dtype == torch.int64
-#     else:
-#         assert "binary" not in out.keys()
-
-#     if return_type:
-#         assert "type" in out.keys()
-#         assert out["type"].dtype == torch.int64
-#     else:
-#         assert "type" not in out.keys()
-
-#     if return_sem:
-#         assert "sem" in out.keys()
-#         assert out["sem"].dtype == torch.int64
-#     else:
-#         assert "sem" not in out.keys()
-
-#     assert "smoothdist" in out.keys()
-#     assert out["image"].dtype == torch.float32
+import pytest
+import torch
+
+from cellseg_models_pytorch.datasets.hdf5_dataset import SegmentationHDF5Dataset
+
+img_transforms = ["rigid", "blur"]
+inst_transforms = ["smooth_dist"]
+
+
+@pytest.mark.optional
+@pytest.mark.parametrize("return_inst", [True, False])
+@pytest.mark.parametrize("return_type", [True, False])
+@pytest.mark.parametrize("return_sem", [True, False])
+@pytest.mark.parametrize("normalization", [None, "minmax"])
+def test_hdf5_dataset(hdf5db, return_inst, return_type, return_sem, normalization):
+    ds = SegmentationHDF5Dataset(
+        path=hdf5db,
+        img_transforms=img_transforms,
+        inst_transforms=inst_transforms,
+        normalization=normalization,
+        return_inst=return_inst,
+        return_type=return_type,
+        return_sem=return_sem,
+    )
+
+    out = next(iter(ds))
+
+    if return_inst:
+        assert "inst" in out.keys()
+        assert out["inst"].dtype == torch.int64
+    else:
+        assert "binary" not in out.keys()
+
+    if return_type:
+        assert "type" in out.keys()
+        assert out["type"].dtype == torch.int64
+    else:
+        assert "type" not in out.keys()
+
+    if return_sem:
+        assert "sem" in out.keys()
+        assert out["sem"].dtype == torch.int64
+    else:
+        assert "sem" not in out.keys()
+
+    assert "smoothdist" in out.keys()
+    assert out["image"].dtype == torch.float32

cellseg_models_pytorch/inference/__init__.py

@@ -1,6 +1,13 @@
+from .folder_dataset import FolderDataset
 from .post_processor import PostProcessor
 from .predictor import Predictor
 from .resize_inferer import ResizeInferer
 from .sliding_window_inferer import SlidingWindowInferer
 
-__all__ = ["Predictor", "PostProcessor", "ResizeInferer", "SlidingWindowInferer"]
+__all__ = [
+    "Predictor",
+    "PostProcessor",
+    "ResizeInferer",
+    "SlidingWindowInferer",
+    "FolderDataset",
+]

cellseg_models_pytorch/inference/_base_inferer.py

@@ -11,9 +11,9 @@
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from ..datasets import FolderDataset
 from ..utils import tensor_to_ndarray
 from ..utils.save_utils import mask2mat
+from .folder_dataset import FolderDataset
 from .post_processor import PostProcessor
 from .predictor import Predictor
 

cellseg_models_pytorch/datasets/folder_dataset.py renamed to cellseg_models_pytorch/inference/folder_dataset.py

@@ -2,22 +2,14 @@
 
 from cellseg_models_pytorch.inference import ResizeInferer, SlidingWindowInferer
 from cellseg_models_pytorch.models import cellpose_plus
-from cellseg_models_pytorch.training import SegmentationExperiment
 
 
 @pytest.mark.parametrize("batch_size", [1, 2])
 def test_slidingwin_inference(img_dir, batch_size):
     model = cellpose_plus(sem_classes=3, type_classes=3, long_skip="unet")
 
-    experiment = SegmentationExperiment(
-        model=model,
-        branch_losses={"cellpose": "mse_ssim", "sem": "ce_dice", "type": "ce_dice"},
-        branch_metrics={"cellpose": [None], "sem": ["miou"], "type": ["miou"]},
-        lookahead=False,
-    )
-
     inferer = SlidingWindowInferer(
-        experiment,
+        model,
         img_dir,
         out_activations={"sem": "softmax", "type": "softmax", "cellpose": "tanh"},
         out_boundary_weights={"sem": False, "type": False, "cellpose": True},
@@ -42,15 +34,8 @@ def test_slidingwin_inference(img_dir, batch_size):
 def test_resize_inference(img_dir, batch_size):
     model = cellpose_plus(sem_classes=3, type_classes=3, long_skip="unet")
 
-    experiment = SegmentationExperiment(
-        model=model,
-        branch_losses={"cellpose": "mse_ssim", "sem": "ce_dice", "type": "ce_dice"},
-        branch_metrics={"cellpose": [None], "sem": ["miou"], "type": ["miou"]},
-        lookahead=False,
-    )
-
     inferer = ResizeInferer(
-        experiment,
+        model,
         img_dir,
         out_activations={"sem": "softmax", "type": "softmax", "cellpose": "tanh"},
         out_boundary_weights={"sem": False, "type": False, "cellpose": True},