Update check scripts

constantinpape · Jan 9, 2025 · e14e240 · e14e240
1 parent 19c8b32
commit e14e240
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/scripts/datasets/histopathology/check_cryonuseg.py b/scripts/datasets/histopathology/check_cryonuseg.py
@@ -15,6 +15,7 @@ def check_cryonuseg():
         path=os.path.join(ROOT, "cryonuseg"),
         patch_shape=(1, 512, 512),
         batch_size=1,
+        split="train",
         rater="b1",
         download=True,
     )

diff --git a/torch_em/data/datasets/histopathology/cryonuseg.py b/torch_em/data/datasets/histopathology/cryonuseg.py
@@ -47,15 +47,19 @@ def _create_split_csv(path, data_dir, split):
     return split_list
 
 
-def get_cryonuseg_data(path: Union[os.PathLike, str], download: bool = False):
+def get_cryonuseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
     """Download the CryoNuSeg dataset for nucleus segmentation.
 
     Args:
         path: Filepath to a folder where the downloaded data will be saved.
         download: Whether to download the data if it is not present.
+
+    Returns:
+        The folder where the data is downloaded and preprocessed.
     """
+    data_dir = os.path.join(path, r"tissue images")
     if os.path.exists(os.path.join(path, r"tissue images")):
-        return
+        return data_dir
 
     os.makedirs(path, exist_ok=True)
     util.download_source_kaggle(
@@ -65,6 +69,8 @@ def get_cryonuseg_data(path: Union[os.PathLike, str], download: bool = False):
     zip_path = os.path.join(path, "segmentation-of-nuclei-in-cryosectioned-he-images.zip")
     util.unzip(zip_path=zip_path, dst=path)
 
+    return data_dir
+
 
 def get_cryonuseg_paths(
     path: Union[os.PathLike, str],
@@ -84,7 +90,7 @@ def get_cryonuseg_paths(
         List of filepaths to the image data.
         List of filepaths to the label data.
     """
-    get_cryonuseg_data(path, download)
+    data_dir = get_cryonuseg_data(path, download)
 
     if rater_choice == "b1":
         label_dir = r"Annotator 1 (biologist)/"
@@ -98,8 +104,12 @@ def get_cryonuseg_paths(
     # Point to the instance labels folder
     label_dir += r"label masks modify"
     split_list = _create_split_csv(path, label_dir, split)
+
+    # Get the raw and label paths
     label_paths = natsorted([os.path.join(path, label_dir, f'{fname}.tif') for fname in split_list])
-    raw_paths = natsorted([os.path.join(path, r"tissue images", f'{fname}.tif') for fname in split_list])
+    raw_paths = natsorted([os.path.join(data_dir, f'{fname}.tif') for fname in split_list])
+
+    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 
     return raw_paths, label_paths