torch_em.data.datasets.light_microscopy.lapd_mouse

  1"""The LAPD Mouse dataset (Lung Anatomy + Particle Deposition mouse archive) contains 3D cryomicrotome
  2fluorescence imaging volumes of 34 mouse lungs with annotations of the airway tree.
  3
  4The raw data is the autofluorescence channel of the imaging cryomicrotome, which shows the anatomical structures
  5(lung, fissures, airway walls). The full resolution volumes have a voxel size of about 9 x 9 x 9.5 um and a size of
  6about 2000 x 1500 x 2500 voxels (more than 10 GB per mouse). The archive also provides versions downsampled by a
  7factor of 2 ('sub2', about 1.5 GB per mouse) and 4 ('sub4', about 200 MB per mouse), which can be chosen via the
  8`resolution` argument.
  9
 10The labels are the airway segment labelmaps: every airway segment (branch) from the trachea to the terminal
 11bronchi has its own id (1 to about 1800 per mouse), which corresponds to the segment ids in the airway tree tables
 12of the archive. Use `labels > 0` to obtain a binary airway mask. The labelmaps are provided at full resolution only
 13and are downsampled by strided subsampling to match the chosen raw resolution.
 14
 15The raw data and labels are converted to a hdf5 file per mouse, with the keys 'raw' (uint16) and 'labels' (uint16).
 16
 17The dataset is located at https://cebs-ext.niehs.nih.gov/cahs/report/lapd/web-download-links.
 18
 19This dataset is from the archive https://doi.org/10.25820/9arg-9w56 (Beichel et al., University of Iowa, 2019).
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24import zlib
 25from typing import Union, Tuple, Literal, List, Optional, Sequence
 26
 27import numpy as np
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36URL = "https://cebs-ext.niehs.nih.gov/cahs/file/download/lapd/{mouse_id}/{filename}"
 37
 38MOUSE_IDS = [f"m{i:02d}" for i in range(1, 35)]
 39
 40RESOLUTIONS = {"full": ("", 1), "sub2": ("Sub2", 2), "sub4": ("Sub4", 4)}
 41
 42# Checksums are only available for the 'sub4' raw volumes and the labels; 'full' and 'sub2' are not verified.
 43CHECKSUMS = {
 44    "m01": {
 45        "sub4": "af22d8ce03860098926bad1130518a3dd612f8d6f471341bd67b381255d6da1c",
 46        "labels": "c12945ee05bd616a58290bdf7f3aee5fe498511ff42a77ebadabff441ebd6666",
 47    },
 48    "m02": {
 49        "sub4": "0563ee1e984d4d4a9fabb49bc8615d8ae302c6422f835b459a36335a73f338a1",
 50        "labels": "82b6d22b49e87bb53b4c82e69a6c0fedf32cb1567ba41857e1eeb945d414f95d",
 51    },
 52    "m03": {
 53        "sub4": "83663a32b770066edd2cfe5524e5b0675f41291c2cb4046a9abf4c426b588898",
 54        "labels": "3195006484650331a19895e41dbec1a0fe0134b4ed5d93a71f61c225a772fe1b",
 55    },
 56    "m04": {
 57        "sub4": "35ded7031395af0b1f65a2b75184173aff95e2e4022df9a7815b4e1fd1c30916",
 58        "labels": "995740d07ef38dded21401ceb4b1d1eed03a6cb0456d3c25471599509e687b71",
 59    },
 60    "m05": {
 61        "sub4": "72736c2caf3f2c88fe0740e6b50d462e383187612f67758d8d4c2d0ccda7b072",
 62        "labels": "b086d2e09f206354578ef19b5d9f4ccd978db2bf68924de07d84a74baf2de036",
 63    },
 64    "m06": {
 65        "sub4": "71827c4d3429892592b3d5b20d96ba236036a87cf358f080efb0e7dd1381056e",
 66        "labels": "8bd0f6b3c982653ed7bac790dc886374b3b8072f2a41a119c5a009447b1842f8",
 67    },
 68    "m07": {
 69        "sub4": "9ae21f57649f6402b5abf7b6262fdf3c2385d95970ff60b7f29f38c5c8e90dc9",
 70        "labels": "33d3d23f026d37cbf5c5f6394f31f3704166b5b56654b5bb4778e54e3bd45504",
 71    },
 72    "m08": {
 73        "sub4": "772d1ecfb5b1c618d5a42c43adef06adaae4f5735a2c443c7314fb9f0f82edb3",
 74        "labels": "00cbec1610f5cf69c3fe46e402eef619e2c3acefc01a5059cb9acd65c8427b70",
 75    },
 76    "m09": {
 77        "sub4": "375243048f41773a034091c10e1804293d25548fb34e8e227783f087f76e329f",
 78        "labels": "ff20215d2c31e615f907e471b15100142139c107eff7715f1b7d32268f37f213",
 79    },
 80    "m10": {
 81        "sub4": "f5d6d149110e4b6561a81f5eb81c274b626bed98cd4283dec6f2009002f10330",
 82        "labels": "13c6d12947a4a315d7cc4e68fe36ae60fcac2eb5c5e2af72337993eba16c10d0",
 83    },
 84    "m11": {
 85        "sub4": "d10a5a5dbc55b97902d5b041866e89d87b39efbb2f6b142b672e97c02734e10a",
 86        "labels": "9c3692147181daa16c50cf0314fec82b658922aca62eebcd10550462f6a8155d",
 87    },
 88    "m12": {
 89        "sub4": "7ff55b74f8c4c46ce09fec86580c44612105d21e8f124eb809c3a11c5f613d5f",
 90        "labels": "4e66ba6cc52d44740c7186ecafba6332d22d03707a14b933535acac4bea91d0f",
 91    },
 92    "m13": {
 93        "sub4": "21689e429107b1927774076a421033bcf372ac7c4c19604b6ae011f7ea9d5bbd",
 94        "labels": "96b826ba0026f24c558b2c044f8b20f81508bbc885be0227b2e1c6809599a89a",
 95    },
 96    "m14": {
 97        "sub4": "ef609e3ab5b93eaad58c73d6b2dd86752fb22dd2c8f33f24f366d7753fab29a2",
 98        "labels": "8bf29236b71572cc3130b71d7576840a1ab99202ee40c590484a01aeccacc2c6",
 99    },
100    "m15": {
101        "sub4": "75a6fc910332a7024ee96ca3ec394ffe75d45b71b0d6660adbb57fad4d3d3132",
102        "labels": "47e6c83231153f48e709761f128f03248baffa096590675451fb75f81f6bd17e",
103    },
104    "m16": {
105        "sub4": "3ff0201879e65573fab8573014c2293606dfaffb9132e56f622925865180999b",
106        "labels": "8027e84785b307624f91d16e03c30321705fe3c5c5615275a7551a345a13a10b",
107    },
108    "m17": {
109        "sub4": "e6e01e40cb924fbc1eb5794dc2b3d6a6e1e906212ae3113aab7dd2e3a678a26f",
110        "labels": "dd32f3429312771d9f75a3072386f5d2147b5eb152f6f6b96e8106359465ff7e",
111    },
112    "m18": {
113        "sub4": "f553f05573f6cdd86fecac1b1c214fe56695460b0b13dae12536a38028fbb747",
114        "labels": "aa06f5e43e903d1c1b33506ec76ac58e12486471d7d7f69bb13289c79df1c1c3",
115    },
116    "m19": {
117        "sub4": "6f9efe9d16387dd37c52c517ff5dc036e551a24ba444df0224035596ac1eabbd",
118        "labels": "36aaf85e14255dbb7273ba110bf536fa0f8fefd3487bfc550482ad98af992e96",
119    },
120    "m20": {
121        "sub4": "23339d32e4b0b6380ea5488e6ddddc90a501080bb605bb87c172aa22427aac9e",
122        "labels": "a916b9131552e0a39e8d7a4ecd4d7996129feca9bd306f4690426e7f71509f2b",
123    },
124    "m21": {
125        "sub4": "d64828d990ef01263aaf88e084693ce52e9011dc71e1931a68978b0f5c5dff33",
126        "labels": "102a53162c20ea8c0f83be4d2fd36682e295caab40cb905cbaaeee6de0966a42",
127    },
128    "m22": {
129        "sub4": "98142d364a38bc98fbc03eadb38d540661e52cd21fb68c32ddbff7c253111d07",
130        "labels": "90319b560b018709edd2b31bc0e1f3b250a2ce9216b797c47b53544bc3983307",
131    },
132    "m23": {
133        "sub4": "ed966b20e7cdb965c86d7d3f072c4ee601dfd6fd876b5eb63f87e910c4237ed9",
134        "labels": "b37f2e9d05755f946c1714bbeab8afbfaa18da27ff0f28eb87f06d413f6b141e",
135    },
136    "m24": {
137        "sub4": "eda4cd76e1d409e6b6da05c70e460a6b0d51858cf5c2c932c415a3ae1cc7b5fc",
138        "labels": "17cc265edd1b795ac200f13caf0391880b45c656141fc93c1bed95bd53f0dae0",
139    },
140    "m25": {
141        "sub4": "3ba9e1151ba7f0f555707160ecbc6070e027efbed51045a1a76f4190619a54ae",
142        "labels": "713d7d81c7d193753465ee664970d7d6c32a7f2d0c0c1421f4a208ac109ac1c1",
143    },
144    "m26": {
145        "sub4": "1c764582d97abdd837d3cceedd3f21d60faefbace9fa77c984dbc1821c4b757b",
146        "labels": "7a9c1e8e745dc3246429a8de7e8cf48b96a9f0d74749b73bf283a9cc7a4a8972",
147    },
148    "m27": {
149        "sub4": "64c1748e992f75caa05caca84c305ffa75d53d22e1a3f64fc55bf0eef66d810c",
150        "labels": "5b5f5a977421195a913a79ab9a0997ba2284bf8d42617e28af284a487d7ed7f0",
151    },
152    "m28": {
153        "sub4": "46cdd907a1eaa46d7992cde53de90c83b682404612919012494b6d171c676fe0",
154        "labels": "e33498ebaab19ea0a845f8bab738816d309d306749f4177a5070c6d1bad51169",
155    },
156    "m29": {
157        "sub4": "7a5d76d054cdf9e2f2c7c32cb6bfa70225ef0c1cff12ec7264650680ad11e196",
158        "labels": "ca2871ba15ac74113ca931b74e7be2d27a7e716ec3f89317afa2c93703f362bb",
159    },
160    "m30": {
161        "sub4": "8776de2ec24b142194867466150d315829825c15181181b679f6b424e0f956ec",
162        "labels": "fc9c4ac97d60e5a58d3dec233a8fbf0aa9b448e5c44c6bd57c6825897bff1357",
163    },
164    "m31": {
165        "sub4": "cb419acb89bf97b4a66f23dcd33d05d94fc8e0bb76a5a7076a8e6801c9fb5870",
166        "labels": "2157a24c9bc44feabcafdf9a53af35a0313d7a49c8fd0bb209571c8760d774ae",
167    },
168    "m32": {
169        "sub4": "72fc16c8b20944a1fc53939a6fb1df7fd15bdbe5ca5620af3419c94ce8bf8ed1",
170        "labels": "7e4f1871e64fe3d6e969eae06b0f491de0e2ee03a58d28ccc255ddce45ab65ba",
171    },
172    "m33": {
173        "sub4": "90f1a99935cc5eb4e6e1d638cf80f00241d58eb9ef3c38185b879ace3db351f3",
174        "labels": "82550b9630263762e2effaedabdef15b687655f059a5c6733706cfd575cf8436",
175    },
176    "m34": {
177        "sub4": "68a24ae5f19b79d052c23d69a84c396ec9a5e7e2430f2fc3d8aa658186003766",
178        "labels": "97077748cd6d34bf4c1c5632ea34107ee744488754d76ddbe1a04fcac671f0e6",
179    },
180}
181
182MHA_DTYPES = {
183    "MET_UCHAR": np.uint8, "MET_CHAR": np.int8, "MET_USHORT": np.uint16, "MET_SHORT": np.int16,
184    "MET_UINT": np.uint32, "MET_INT": np.int32, "MET_FLOAT": np.float32, "MET_DOUBLE": np.float64,
185}
186
187
188def _read_mha(path):
189    """Read a MetaImage (.mha) volume with the data stored inside the file. Returns the data in zyx order."""
190    header = {}
191    with open(path, "rb") as f:
192        while True:
193            line = f.readline().decode("ascii")
194            if "=" not in line:
195                continue
196            key, value = [s.strip() for s in line.split("=", 1)]
197            header[key] = value
198            if key == "ElementDataFile":
199                data = f.read()
200                break
201
202    if header["ElementDataFile"] != "LOCAL":
203        raise RuntimeError(f"Only .mha files with local data are supported, got '{header['ElementDataFile']}'.")
204    if header.get("CompressedData", "False") == "True":
205        data = zlib.decompress(data)
206
207    dtype = np.dtype(MHA_DTYPES[header["ElementType"]])
208    if header.get("BinaryDataByteOrderMSB", "False") == "True":
209        dtype = dtype.newbyteorder(">")
210    shape = tuple(int(s) for s in header["DimSize"].split())[::-1]  # DimSize is in xyz order.
211    return np.frombuffer(data, dtype=dtype).reshape(shape)
212
213
214def _convert_to_hdf5(raw_path, label_path, output_path, factor):
215    import h5py
216    import nrrd
217
218    raw = _read_mha(raw_path)
219
220    # nrrd returns the data in xyz order, we transpose it to zyx.
221    labels = nrrd.read(label_path)[0].transpose(2, 1, 0)
222    if factor > 1:
223        # The downsampled raw volumes are block averages, so the voxel centers are shifted by (factor - 1) / 2
224        # full resolution voxels. We subsample the labels with a corresponding offset and crop to the raw shape.
225        offset = (factor - 1) // 2
226        labels = labels[offset::factor, offset::factor, offset::factor]
227        labels = labels[:raw.shape[0], :raw.shape[1], :raw.shape[2]]
228
229    if raw.shape != labels.shape:
230        raise RuntimeError(f"Shape mismatch between raw {raw.shape} and labels {labels.shape} for '{raw_path}'.")
231
232    chunks = (1,) + raw.shape[1:]
233    with h5py.File(output_path, "w") as f:
234        f.create_dataset("raw", data=raw, compression="gzip", chunks=chunks)
235        f.create_dataset("labels", data=labels, compression="gzip", chunks=chunks)
236
237
238def get_lapd_mouse_data(
239    path: Union[os.PathLike, str],
240    mouse_id: str,
241    resolution: Literal["full", "sub2", "sub4"] = "sub4",
242    download: bool = False,
243) -> str:
244    """Download and preprocess one mouse of the LAPD Mouse dataset.
245
246    Args:
247        path: Filepath to a folder where the data is downloaded for further processing.
248        mouse_id: The mouse to download. One of the ids in `MOUSE_IDS`.
249        resolution: The resolution of the raw data. Either 'full', 'sub2' (downsampled by 2) or 'sub4'
250            (downsampled by 4).
251        download: Whether to download the data if it is not present.
252
253    Returns:
254        Filepath to the hdf5 file with the raw data and airway labels of this mouse.
255    """
256    if mouse_id not in MOUSE_IDS:
257        raise ValueError(f"'{mouse_id}' is not a valid mouse id. Choose one of {MOUSE_IDS}.")
258    if resolution not in RESOLUTIONS:
259        raise ValueError(f"'{resolution}' is not a valid resolution. Choose one of {list(RESOLUTIONS.keys())}.")
260
261    volume_path = os.path.join(path, f"{mouse_id}_{resolution}.h5")
262    if os.path.exists(volume_path):
263        return volume_path
264
265    download_dir = os.path.join(path, "downloads")
266    os.makedirs(download_dir, exist_ok=True)
267
268    suffix, factor = RESOLUTIONS[resolution]
269    raw_name, label_name = f"{mouse_id}_Autofluorescent{suffix}.mha", f"{mouse_id}_AirwaySegments.nrrd"
270    raw_path, label_path = os.path.join(download_dir, raw_name), os.path.join(download_dir, label_name)
271
272    util.download_source(
273        path=raw_path, url=URL.format(mouse_id=mouse_id, filename=raw_name), download=download,
274        checksum=CHECKSUMS[mouse_id].get(resolution),
275    )
276    util.download_source(
277        path=label_path, url=URL.format(mouse_id=mouse_id, filename=label_name), download=download,
278        checksum=CHECKSUMS[mouse_id]["labels"],
279    )
280
281    _convert_to_hdf5(raw_path, label_path, volume_path, factor)
282    os.remove(raw_path)
283    os.remove(label_path)
284
285    return volume_path
286
287
288def get_lapd_mouse_paths(
289    path: Union[os.PathLike, str],
290    mouse_ids: Optional[Sequence[str]] = None,
291    resolution: Literal["full", "sub2", "sub4"] = "sub4",
292    download: bool = False,
293) -> List[str]:
294    """Get paths to the LAPD Mouse data.
295
296    Args:
297        path: Filepath to a folder where the data is downloaded for further processing.
298        mouse_ids: The mice to use. By default, all 34 mice are used.
299        resolution: The resolution of the raw data. Either 'full', 'sub2' (downsampled by 2) or 'sub4'
300            (downsampled by 4).
301        download: Whether to download the data if it is not present.
302
303    Returns:
304        List of filepaths for the hdf5 volumes, which contain the raw data and airway labels.
305    """
306    if mouse_ids is None:
307        mouse_ids = MOUSE_IDS
308    elif isinstance(mouse_ids, str):
309        mouse_ids = [mouse_ids]
310
311    volume_paths = [get_lapd_mouse_data(path, mouse_id, resolution, download) for mouse_id in mouse_ids]
312    return volume_paths
313
314
315def get_lapd_mouse_dataset(
316    path: Union[os.PathLike, str],
317    patch_shape: Tuple[int, ...],
318    mouse_ids: Optional[Sequence[str]] = None,
319    resolution: Literal["full", "sub2", "sub4"] = "sub4",
320    resize_inputs: bool = False,
321    download: bool = False,
322    **kwargs
323) -> Dataset:
324    """Get the LAPD Mouse dataset for airway segmentation in cryomicrotome fluorescence volumes of mouse lungs.
325
326    Args:
327        path: Filepath to a folder where the data is downloaded for further processing.
328        patch_shape: The patch shape to use for training.
329        mouse_ids: The mice to use. By default, all 34 mice are used.
330        resolution: The resolution of the raw data. Either 'full', 'sub2' (downsampled by 2) or 'sub4'
331            (downsampled by 4).
332        resize_inputs: Whether to resize inputs to the desired patch shape.
333        download: Whether to download the data if it is not present.
334        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
335
336    Returns:
337        The segmentation dataset.
338    """
339    volume_paths = get_lapd_mouse_paths(path, mouse_ids, resolution, download)
340
341    if resize_inputs:
342        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
343        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
344            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
345        )
346
347    return torch_em.default_segmentation_dataset(
348        raw_paths=volume_paths,
349        raw_key="raw",
350        label_paths=volume_paths,
351        label_key="labels",
352        patch_shape=patch_shape,
353        is_seg_dataset=True,
354        **kwargs
355    )
356
357
358def get_lapd_mouse_loader(
359    path: Union[os.PathLike, str],
360    batch_size: int,
361    patch_shape: Tuple[int, ...],
362    mouse_ids: Optional[Sequence[str]] = None,
363    resolution: Literal["full", "sub2", "sub4"] = "sub4",
364    resize_inputs: bool = False,
365    download: bool = False,
366    **kwargs
367) -> DataLoader:
368    """Get the LAPD Mouse dataloader for airway segmentation in cryomicrotome fluorescence volumes of mouse lungs.
369
370    Args:
371        path: Filepath to a folder where the data is downloaded for further processing.
372        batch_size: The batch size for training.
373        patch_shape: The patch shape to use for training.
374        mouse_ids: The mice to use. By default, all 34 mice are used.
375        resolution: The resolution of the raw data. Either 'full', 'sub2' (downsampled by 2) or 'sub4'
376            (downsampled by 4).
377        resize_inputs: Whether to resize inputs to the desired patch shape.
378        download: Whether to download the data if it is not present.
379        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
380
381    Returns:
382        The DataLoader.
383    """
384    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
385    dataset = get_lapd_mouse_dataset(path, patch_shape, mouse_ids, resolution, resize_inputs, download, **ds_kwargs)
386    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)