Minor updates to LM datasets (#483)

constantinpape · Feb 9, 2025 · 89652c9 · 89652c9
1 parent 6a042fa
commit 89652c9
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/torch_em/data/datasets/light_microscopy/arvidsson.py b/torch_em/data/datasets/light_microscopy/arvidsson.py
@@ -62,6 +62,8 @@ def get_arvidsson_data(
     if os.path.exists(data_dir):
         return data_dir
 
+    os.makedirs(path, exist_ok=True)
+
     zip_path = os.path.join(path, f"{dname}.zip")
     util.download_source(path=zip_path, url=URLS[split], download=download, checksum=CHECKSUMS[split])
     util.unzip(zip_path=os.path.join(path, f"{dname}.zip"), dst=path)

diff --git a/torch_em/data/datasets/light_microscopy/bitdepth_nucseg.py b/torch_em/data/datasets/light_microscopy/bitdepth_nucseg.py
@@ -67,7 +67,7 @@ def get_bitdepth_nucseg_data(path: Union[os.PathLike, str], download: bool = Fal
 
 def get_bitdepth_nucseg_paths(
     path: Union[os.PathLike, str],
-    magnification: Optional[Literal['20x', '40x_air', '40x_oil' '63x_oil']] = None,
+    magnification: Optional[Literal['20x', '40x_air', '40x_oil', '63x_oil']] = None,
     download: bool = False
 ) -> Tuple[List[str], List[str]]:
     """Get paths to the BitDepth NucSeg data.
@@ -101,7 +101,7 @@ def get_bitdepth_nucseg_paths(
 def get_bitdepth_nucseg_dataset(
     path: Union[os.PathLike, str],
     patch_shape: Tuple[int, int],
-    magnification: Optional[Literal['20x', '40x_air', '40x_oil' '63x_oil']] = None,
+    magnification: Optional[Literal['20x', '40x_air', '40x_oil', '63x_oil']] = None,
     download: bool = False,
     **kwargs
 ) -> Dataset:
@@ -134,7 +134,7 @@ def get_bitdepth_nucseg_loader(
     path: Union[os.PathLike, str],
     batch_size: int,
     patch_shape: Tuple[int, int],
-    magnification: Optional[Literal['20x', '40x_air', '40x_oil' '63x_oil']] = None,
+    magnification: Optional[Literal['20x', '40x_air', '40x_oil', '63x_oil']] = None,
     download: bool = False,
     **kwargs
 ) -> DataLoader:

diff --git a/torch_em/data/datasets/light_microscopy/brifiseg.py b/torch_em/data/datasets/light_microscopy/brifiseg.py
@@ -63,10 +63,10 @@ def get_brifiseg_data(path: Union[os.PathLike, str], download: bool = False) ->
     if os.path.exists(data_dir):
         return data_dir
 
-    tar_path = os.path.join(path, "1channel.tar")
-    if not os.path.exists(tar_path) and not download:
-        raise ValueError("The tar file is not found and download is set to 'False'.")
+    os.makedirs(path, exist_ok=True)
 
+    tar_path = os.path.join(path, "1channel.tar")
+    util.download_source(path=tar_path, url=URL, checksum=CHECKSUM, download=download)
     util.unzip_tarfile(tar_path=tar_path, dst=path)
 
     for zip_path in glob(os.path.join(os.path.join(path, "1channel"), "*.zip")):