torch_em.data.datasets.light_microscopy.lungvis

  1"""The LungVis 1.0 dataset contains light sheet fluorescence microscopy volumes of tissue-cleared mouse lungs
  2with annotations of the airway tree.
  3
  4The dataset consists of 78 lungs: 37 lungs with nanoparticle delivery via different routes
  5(intranasal liquid aspiration 'INLA', intratracheal liquid instillation 'ITLI', nose-only aerosol inhalation 'NOAI'
  6and ventilator-assisted aerosol delivery 'VAAD') and 41 additional lungs imaged only for the airway geometry
  7('Lung_001' to 'Lung_041'). The raw data is the tissue autofluorescence channel (470 to 570 nm).
  8The airway annotations are of three kinds:
  9- 'MS': manually segmented airways (3 lungs: 'INLA_001', 'Lung_001', 'Lung_002').
 10- 'MCAI': manually corrected nnU-Net predictions (17 lungs).
 11- 'AI': uncorrected nnU-Net predictions (58 lungs).
 12Use `annotation="manual"` (default) to get the 20 lungs with human-verified airway annotations ('MS' and 'MCAI'),
 13`annotation="ai"` for the 58 lungs with automatic annotations and `annotation="all"` for all 78 lungs.
 14
 15Every lung is downloaded as a separate archive (25 MB to 8.4 GB, 100 GB in total). The raw data and the airway
 16labels are converted to a hdf5 file per lung, with keys 'raw' (uint8) and 'labels' (uint8, 1 = airway).
 17Voxel sizes are 5.159 x 5.159 x (10, 15 or 20) um for most lungs.
 18
 19NOTE: 'Lung_002' is a small region of 'Lung_001', 'Lung_010' is the same lung as 'ITLI_001'
 20and 'Lung_013' is the same lung as 'VAAD_001'.
 21
 22The dataset is located at https://zenodo.org/records/7413818.
 23
 24This dataset is from the publication https://doi.org/10.1038/s41467-024-54267-1.
 25Please cite it if you use this dataset in your research.
 26"""
 27
 28import os
 29import re
 30import shutil
 31from glob import glob
 32from natsort import natsorted
 33from typing import Union, Tuple, Literal, List, Optional, Sequence
 34
 35from torch.utils.data import Dataset, DataLoader
 36
 37import torch_em
 38
 39from .. import util
 40
 41
 42URL = "https://zenodo.org/records/7413818/files/{lung_id}.zip"
 43
 44ANNOTATION_TYPES = {
 45    "INLA_001": "MS", "INLA_002": "AI",
 46    "ITLI_001": "AI", "ITLI_002": "MCAI", "ITLI_003": "MCAI", "ITLI_004": "AI", "ITLI_005": "AI", "ITLI_006": "AI",
 47    "ITLI_007": "AI", "ITLI_008": "AI", "ITLI_009": "AI", "ITLI_010": "AI", "ITLI_011": "MCAI", "ITLI_012": "AI",
 48    "ITLI_013": "AI", "ITLI_014": "AI",
 49    "NOAI_001": "MCAI", "NOAI_002": "AI", "NOAI_003": "AI",
 50    "VAAD_001": "AI", "VAAD_002": "MCAI", "VAAD_003": "AI", "VAAD_004": "MCAI", "VAAD_005": "AI", "VAAD_006": "AI",
 51    "VAAD_007": "AI", "VAAD_008": "AI", "VAAD_009": "AI", "VAAD_010": "MCAI", "VAAD_011": "AI", "VAAD_012": "AI",
 52    "VAAD_013": "AI", "VAAD_014": "AI", "VAAD_015": "MCAI", "VAAD_016": "AI", "VAAD_017": "AI", "VAAD_018": "MCAI",
 53    "Lung_001": "MS", "Lung_002": "MS", "Lung_003": "MCAI", "Lung_004": "MCAI", "Lung_005": "MCAI",
 54    "Lung_006": "MCAI", "Lung_007": "MCAI", "Lung_008": "MCAI", "Lung_009": "MCAI", "Lung_010": "MCAI",
 55    "Lung_011": "AI", "Lung_012": "AI", "Lung_013": "AI", "Lung_014": "AI", "Lung_015": "AI", "Lung_016": "AI",
 56    "Lung_017": "AI", "Lung_018": "AI", "Lung_019": "AI", "Lung_020": "AI", "Lung_021": "AI", "Lung_022": "AI",
 57    "Lung_023": "AI", "Lung_024": "AI", "Lung_025": "AI", "Lung_026": "AI", "Lung_027": "AI", "Lung_028": "AI",
 58    "Lung_029": "AI", "Lung_030": "AI", "Lung_031": "AI", "Lung_032": "AI", "Lung_033": "AI", "Lung_034": "AI",
 59    "Lung_035": "AI", "Lung_036": "AI", "Lung_037": "AI", "Lung_038": "AI", "Lung_039": "AI", "Lung_040": "AI",
 60    "Lung_041": "AI",
 61}
 62
 63LUNG_IDS = list(ANNOTATION_TYPES.keys())
 64
 65# Checksums are only available for the archives with manual annotations ('MS' and 'MCAI').
 66CHECKSUMS = {
 67    "INLA_001": "5b23eb6b3d20435606352c50d8d578e6eea88fdfaeea83eb4adcc0fa6180d8a1",
 68    "ITLI_002": "76444796540ab3ce33d52961944c9104925e3697cc4df4a8de84519c4619433a",
 69    "ITLI_003": "c331642a781f91dd258e7f7470f5c554e82c8bdad7ff13e36fecebcd38ae840b",
 70    "ITLI_011": "5ce194f22e963cfd9bfa2251ad3240394ee7af8a9ba9931ba16a39b2d2cbd882",
 71    "NOAI_001": "9b3aa922cee98541c98430a137b95bcc28deb937336613407cbd40ef556ae257",
 72    "VAAD_002": "ea007e5f41644d2a3d6cc19fe9f1ec48c66d618bd4b0ad3676b36683be3ac7e6",
 73    "VAAD_004": "bff40988bcd2932ebb0579d77b641c385c79f124ad7c9d819436568e105bccc3",
 74    "VAAD_010": "51f3d0fff0bfb509dab3646a3fe1c961b36787e2d47436a2d5339880d7d3d421",
 75    "VAAD_015": "359294fb0ab4a42a6e2935255a54368c4a1b874d913f4d7e5309d27593d27589",
 76    "VAAD_018": "c712a57589ea467e7c066ee2c4aeae7e71baeac9b4c3792f7f56197fb0504ea8",
 77    "Lung_001": "589eee28e1ac75998a3fa6baf32e164d94ec58ff61aee1edeae0e3e932bd5879",
 78    "Lung_002": "783b794dfe98407330923b904d07c5375e8931ca7c534c07821168e1e9866749",
 79    "Lung_003": "80943302fcac472fd5a28e849575c843819250c0e841941ac927f5574e0c5b74",
 80    "Lung_004": "41c3dea4bc9e06951af778198b36d3a6ac0163c543c40ec825e464270f401b9b",
 81    "Lung_005": "96d16eabf3ef0060402172f3853c337a709e9c44629f221c8c8890f00c3d13ca",
 82    "Lung_006": "06bc7f382272362c4fea323b5b0eed730103602b383e97096df8fb915edfd397",
 83    "Lung_007": "ff34057643fbf7d38d669489be86f1322dff140c55ef269ebf44133ce3ff137d",
 84    "Lung_008": "de4d68f8ebe32837a943a09306384ab9b49c8a5cc3c4b1e180b428561aee171a",
 85    "Lung_009": "626fbc9d07d6c080f95f9647abd29bf9a2fd2fd9e0b03f78daa9f370776acfee",
 86    "Lung_010": "02a57bdcefc5c7de5f5ea49e690435e464f722f1792b301bfd97b02c0192dac4",
 87}
 88
 89
 90def _find_volumes(extracted_dir, lung_id):
 91    tif_paths = natsorted(glob(os.path.join(extracted_dir, "**", "*.tif"), recursive=True))
 92    tif_names = {os.path.basename(p): p for p in tif_paths}
 93
 94    # The labels are the MS, MCAI or AI airway volumes, e.g. 'ITLI_003_MCAI_airway AF545nm_UID_11-30-39.tif',
 95    # 'Lung_003_MCAI results_UID_19-49-29.tif' or 'ITLI_002_MCAI_AF545nm_UID_10-28-58.tif'.
 96    label_names = [name for name in tif_names if re.search(r"_(MS|MCAI|AI)[ _]", name)]
 97    if len(label_names) != 1:
 98        raise RuntimeError(f"Expected exactly one label volume for '{lung_id}', found {label_names}.")
 99    label_name = label_names[0]
100
101    # The raw data is the tissue autofluorescence channel, e.g. 'Lung_002_Raw AF545nm_UID_11-08-31.tif'.
102    # Some lungs have multiple autofluorescence channels. We pick the one the airways were annotated on if the
103    # label name contains the wavelength, otherwise the channel with the shortest wavelength (the primary AF channel).
104    raw_names = {}
105    for name in tif_names:
106        match = re.search(r"[Rr]aw[ _]?AF(\d+)nm", name)
107        if match is not None:
108            raw_names[int(match.group(1))] = name
109    if len(raw_names) == 0:
110        raise RuntimeError(f"Could not find the raw volume for '{lung_id}' in {list(tif_names.keys())}.")
111    annotated = [wavelength for wavelength in raw_names if str(wavelength) in label_name]
112    raw_name = raw_names[annotated[0] if len(annotated) == 1 else min(raw_names)]
113
114    return tif_names[raw_name], tif_names[label_name]
115
116
117def _convert_to_hdf5(extracted_dir, output_path, lung_id):
118    import h5py
119    import tifffile
120
121    raw_path, label_path = _find_volumes(extracted_dir, lung_id)
122
123    raw = tifffile.imread(raw_path)
124    labels = tifffile.imread(label_path)
125    if raw.shape != labels.shape:
126        raise RuntimeError(f"Shape mismatch for '{lung_id}': raw {raw.shape} vs labels {labels.shape}.")
127
128    # The airway foreground is stored either as 1 or as 255. We normalize it to 1.
129    labels = (labels > 0).astype("uint8")
130
131    with h5py.File(output_path, "w") as f:
132        f.create_dataset("raw", data=raw, compression="gzip", chunks=(1,) + raw.shape[1:])
133        f.create_dataset("labels", data=labels, compression="gzip", chunks=(1,) + labels.shape[1:])
134
135
136def get_lungvis_data(path: Union[os.PathLike, str], lung_id: str, download: bool = False) -> str:
137    """Download and preprocess one lung of the LungVis dataset.
138
139    Args:
140        path: Filepath to a folder where the data is downloaded for further processing.
141        lung_id: The lung to download. One of the ids in `LUNG_IDS`.
142        download: Whether to download the data if it is not present.
143
144    Returns:
145        Filepath to the hdf5 file with the raw data and airway labels of this lung.
146    """
147    if lung_id not in LUNG_IDS:
148        raise ValueError(f"'{lung_id}' is not a valid lung id. Choose one of {LUNG_IDS}.")
149
150    volume_path = os.path.join(path, f"{lung_id}.h5")
151    if os.path.exists(volume_path):
152        return volume_path
153
154    os.makedirs(path, exist_ok=True)
155
156    zip_path = os.path.join(path, f"{lung_id}.zip")
157    util.download_source(
158        path=zip_path, url=URL.format(lung_id=lung_id), download=download, checksum=CHECKSUMS.get(lung_id)
159    )
160    extracted_dir = os.path.join(path, lung_id)
161    util.unzip(zip_path=zip_path, dst=path)
162
163    _convert_to_hdf5(extracted_dir, volume_path, lung_id)
164    shutil.rmtree(extracted_dir)
165
166    return volume_path
167
168
169def get_lungvis_paths(
170    path: Union[os.PathLike, str],
171    annotation: Literal["manual", "ai", "all"] = "manual",
172    lung_ids: Optional[Sequence[str]] = None,
173    download: bool = False,
174) -> List[str]:
175    """Get paths to the LungVis data.
176
177    Args:
178        path: Filepath to a folder where the data is downloaded for further processing.
179        annotation: The type of airway annotations. Either 'manual' (manually segmented or manually corrected),
180            'ai' (uncorrected nnU-Net predictions) or 'all'.
181        lung_ids: The lungs to use. By default, all lungs matching the chosen annotation type are used.
182        download: Whether to download the data if it is not present.
183
184    Returns:
185        List of filepaths for the hdf5 volumes, which contain the raw data and airway labels.
186    """
187    if annotation not in ("manual", "ai", "all"):
188        raise ValueError(f"'{annotation}' is not a valid annotation type. Choose one of 'manual', 'ai' or 'all'.")
189
190    if lung_ids is None:
191        if annotation == "manual":
192            lung_ids = [lid for lid, atype in ANNOTATION_TYPES.items() if atype in ("MS", "MCAI")]
193        elif annotation == "ai":
194            lung_ids = [lid for lid, atype in ANNOTATION_TYPES.items() if atype == "AI"]
195        else:
196            lung_ids = LUNG_IDS
197    elif isinstance(lung_ids, str):
198        lung_ids = [lung_ids]
199
200    volume_paths = [get_lungvis_data(path, lung_id, download) for lung_id in lung_ids]
201    return volume_paths
202
203
204def get_lungvis_dataset(
205    path: Union[os.PathLike, str],
206    patch_shape: Tuple[int, ...],
207    annotation: Literal["manual", "ai", "all"] = "manual",
208    lung_ids: Optional[Sequence[str]] = None,
209    resize_inputs: bool = False,
210    download: bool = False,
211    **kwargs
212) -> Dataset:
213    """Get the LungVis dataset for airway segmentation in light sheet microscopy volumes of mouse lungs.
214
215    Args:
216        path: Filepath to a folder where the data is downloaded for further processing.
217        patch_shape: The patch shape to use for training.
218        annotation: The type of airway annotations. Either 'manual' (manually segmented or manually corrected),
219            'ai' (uncorrected nnU-Net predictions) or 'all'.
220        lung_ids: The lungs to use. By default, all lungs matching the chosen annotation type are used.
221        resize_inputs: Whether to resize inputs to the desired patch shape.
222        download: Whether to download the data if it is not present.
223        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
224
225    Returns:
226        The segmentation dataset.
227    """
228    volume_paths = get_lungvis_paths(path, annotation, lung_ids, download)
229
230    if resize_inputs:
231        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
232        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
233            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
234        )
235
236    return torch_em.default_segmentation_dataset(
237        raw_paths=volume_paths,
238        raw_key="raw",
239        label_paths=volume_paths,
240        label_key="labels",
241        patch_shape=patch_shape,
242        is_seg_dataset=True,
243        **kwargs
244    )
245
246
247def get_lungvis_loader(
248    path: Union[os.PathLike, str],
249    batch_size: int,
250    patch_shape: Tuple[int, ...],
251    annotation: Literal["manual", "ai", "all"] = "manual",
252    lung_ids: Optional[Sequence[str]] = None,
253    resize_inputs: bool = False,
254    download: bool = False,
255    **kwargs
256) -> DataLoader:
257    """Get the LungVis dataloader for airway segmentation in light sheet microscopy volumes of mouse lungs.
258
259    Args:
260        path: Filepath to a folder where the data is downloaded for further processing.
261        batch_size: The batch size for training.
262        patch_shape: The patch shape to use for training.
263        annotation: The type of airway annotations. Either 'manual' (manually segmented or manually corrected),
264            'ai' (uncorrected nnU-Net predictions) or 'all'.
265        lung_ids: The lungs to use. By default, all lungs matching the chosen annotation type are used.
266        resize_inputs: Whether to resize inputs to the desired patch shape.
267        download: Whether to download the data if it is not present.
268        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
269
270    Returns:
271        The DataLoader.
272    """
273    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
274    dataset = get_lungvis_dataset(path, patch_shape, annotation, lung_ids, resize_inputs, download, **ds_kwargs)
275    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)