-
Notifications
You must be signed in to change notification settings - Fork 19
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
37c9719
commit c0bf70c
Showing
2 changed files
with
98 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,97 @@ | ||
"""This dataset contains annotation for nucleus segmentation in 3d fluorescence microscopy. | ||
This dataset is from the publication https://doi.org/10.1101/2024.05.17.594691 . | ||
Please cite it if you use this dataset in your research. | ||
""" | ||
|
||
import os | ||
from glob import glob | ||
from typing import Tuple, Union | ||
|
||
import torch_em | ||
from torch.utils.data import Dataset, DataLoader | ||
from .. import util | ||
|
||
URL = "https://zenodo.org/records/11095111/files/DATASET_WITH_GT.zip?download=1" | ||
CHECKSUM = "6d8e8d778e479000161fdfea70201a6ded95b3958a703f69def63e69bbddf9d6" | ||
|
||
|
||
def get_cellseg_3d_data(path: Union[os.PathLike, str], download: bool) -> str: | ||
"""Download the CellSeg3d training data. | ||
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 filepath to the training data. | ||
""" | ||
url = URL | ||
checksum = CHECKSUM | ||
|
||
data_path = os.path.join(path, "DATASET_WITH_GT") | ||
if os.path.exists(data_path): | ||
return data_path | ||
|
||
os.makedirs(path, exist_ok=True) | ||
zip_path = os.path.join(path, "cellseg3d.zip") | ||
util.download_source(zip_path, url, download, checksum) | ||
util.unzip(zip_path, path, True) | ||
|
||
return data_path | ||
|
||
|
||
def get_cellseg_3d_dataset( | ||
path: Union[os.PathLike, str], | ||
patch_shape: Tuple[int, int], | ||
download: bool = False, | ||
**kwargs | ||
) -> Dataset: | ||
"""Get the CellSeg3d dataset for segmenting nuclei in 3d fluorescence microscopy. | ||
Args: | ||
path: Filepath to a folder where the downloaded data will be saved. | ||
patch_shape: The patch shape to use for training. | ||
download: Whether to download the data if it is not present. | ||
kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`. | ||
Returns: | ||
The segmentation dataset. | ||
""" | ||
data_root = get_cellseg_3d_data(path, download) | ||
|
||
raw_paths = sorted(glob(os.path.join(data_root, "*.tif"))) | ||
label_paths = sorted(glob(os.path.join(data_root, "labels", "*.tif"))) | ||
|
||
raw_key, label_key = None, None | ||
|
||
return torch_em.default_segmentation_dataset( | ||
raw_paths, raw_key, label_paths, label_key, patch_shape, **kwargs | ||
) | ||
|
||
|
||
def get_cellseg_3d_loader( | ||
path: Union[os.PathLike, str], | ||
patch_shape: Tuple[int, int], | ||
batch_size: int, | ||
download: bool = False, | ||
**kwargs | ||
) -> DataLoader: | ||
"""Get the CovidIF dataloder for segmenting nuclei or cells in immunofluorescence microscopy. | ||
Args: | ||
path: Filepath to a folder where the downloaded data will be saved. | ||
patch_shape: The patch shape to use for training. | ||
batch_size: The batch size for training. | ||
download: Whether to download the data if it is not present. | ||
kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader. | ||
Returns: | ||
The DataLoader. | ||
""" | ||
ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs) | ||
dataset = get_cellseg_3d_dataset( | ||
path, patch_shape, download=download, **ds_kwargs, | ||
) | ||
loader = torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs) | ||
return loader |