torch_em.data.datasets.light_microscopy.bbbc046

  1"""The BBBC046 dataset (FiloData3D) contains synthetic 3D time-lapse fluorescence microscopy images of
  2single A549 lung cancer cells with filopodia, and ground truth masks of the cell body and filopodia.
  3
  4The dataset consists of 180 synthetic sequences of 30 frames each (5400 volumes). They are derived from
  59 base sequences, 3 per cell phenotype: wild-type ('WT-ID550', 'WT-ID551', 'WT-ID552'),
  6CRMP-2-overexpressing ('OE-ID350', 'OE-ID351', 'OE-ID352') and CRMP-2-phospho-defective
  7('PD-ID450', 'PD-ID451', 'PD-ID452'). Every base sequence is rendered for 4 anisotropy ratios
  8(1, 2, 4 and 8, which reduce the number of z-slices) and 5 fluorescence level factors
  9(0.25, 0.50, 1.00, 2.00 and 4.00, which change the signal-to-noise ratio).
 10The ground truth masks are shared across the fluorescence level factors of a sequence.
 11
 12The label ids are: 0 (background), 50 (cell body) and ids of 100 and above for the filopodia. The k-th filopodium
 13uses the ids 100 * k (its primary branch) and 100 * k + 1, 100 * k + 2, ... (its side branches). These ids are also
 14used in the accompanying trajectory and length text files.
 15
 16The base sequences are downloaded separately, each one as an archive of 1.2 to 9.8 GB (36 GB in total).
 17
 18The dataset is located at https://bbbc.broadinstitute.org/BBBC046.
 19
 20This dataset is from the publications https://doi.org/10.1109/TMI.2018.2845884 and
 21https://doi.org/10.1109/ICIP.2019.8803721.
 22Please cite them if you use this dataset in your research.
 23"""
 24
 25import os
 26import shutil
 27from glob import glob
 28from natsort import natsorted
 29from typing import Union, Tuple, Literal, List, Optional, Sequence
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38SEQUENCE_IDS = [
 39    "OE-ID350", "OE-ID351", "OE-ID352", "PD-ID450", "PD-ID451", "PD-ID452", "WT-ID550", "WT-ID551", "WT-ID552",
 40]
 41
 42URLS = {sequence_id: f"https://data.broadinstitute.org/bbbc/BBBC046/{sequence_id}.zip" for sequence_id in SEQUENCE_IDS}
 43
 44CHECKSUMS = {
 45    "OE-ID350": "ee22ecdb33311be4533b00bae4866604e1062466056a46a575486bdfefe2c487",
 46    "OE-ID351": "da496e9018c584ea15a7b12fca6e625b448f5acb9b844782755dadd443000a95",
 47    "OE-ID352": "cb0ae686a980abb2318e03513420ed007f0b30044079778832809b64dd98f4eb",
 48    "PD-ID450": "b67cc148bf38674e3cbecb4f7b044dec3a1b32a3363c53d14baeab7fbb88a0dc",
 49    "PD-ID451": "ce39ad00131beda3f41e886377ac2030d866213ff8b39f302b15c493da5086c0",
 50    "PD-ID452": "3f0fb99a96af7299f9029a93ec867be0e94cd18daae8919a40cfa9cfaa82a924",
 51    "WT-ID550": "2dcf6dbd1faff96fb919aeb7249cf1f25196729ff8304f556d5275a223af878b",
 52    "WT-ID551": "c6e3216bc50ce76d66bd4a35e2d452528aa939852f6f513cf3586dc8e2603e4f",
 53    "WT-ID552": "222c35d171f9c3f1808f302bc9f53700af99267bfede5f54701918d929559f7a",
 54}
 55
 56ANISOTROPY_RATIOS = [1, 2, 4, 8]
 57FLUORESCENCE_FACTORS = ["0.25", "0.50", "1.00", "2.00", "4.00"]
 58
 59
 60def _unzip_with_offset_fix(zip_path, dst):
 61    """Extract a zip archive whose central directory stores wrong local header offsets.
 62
 63    The archives of the 'PD' sequences (larger than 4 GB) were created with a tool that wrote local header offsets
 64    shifted by a multiple of 4 GB, which `zipfile` cannot read directly. We correct the offsets before extraction.
 65    """
 66    import zipfile
 67
 68    with zipfile.ZipFile(zip_path) as f:
 69        infos = f.infolist()
 70        for info in infos:
 71            for shift in (0, -2**32, 2**32, -2**33, 2**33):
 72                offset = info.header_offset + shift
 73                if offset < 0:
 74                    continue
 75                f.fp.seek(offset)
 76                if f.fp.read(4) == b"PK\x03\x04":
 77                    info.header_offset = offset
 78                    break
 79            else:
 80                raise RuntimeError(f"Could not locate the local header of '{info.filename}' in '{zip_path}'.")
 81
 82        # Recent python versions check for overlapping entries with the end offsets computed when opening
 83        # the archive. We recompute them for the corrected header offsets.
 84        infos = sorted(infos, key=lambda info: info.header_offset)
 85        for info, next_info in zip(infos, infos[1:] + [None]):
 86            if hasattr(info, "_end_offset"):
 87                info._end_offset = f.start_dir if next_info is None else next_info.header_offset
 88
 89        for info in infos:
 90            f.extract(info, dst)
 91
 92    os.remove(zip_path)
 93
 94
 95def get_bbbc046_data(path: Union[os.PathLike, str], sequence_id: str, download: bool = False) -> List[str]:
 96    """Download one base sequence of the BBBC046 dataset.
 97
 98    Args:
 99        path: Filepath to a folder where the data is downloaded for further processing.
100        sequence_id: The base sequence to download. One of the ids in `SEQUENCE_IDS`.
101        download: Whether to download the data if it is not present.
102
103    Returns:
104        List of filepaths to the sequence folders, one per anisotropy ratio.
105    """
106    if sequence_id not in SEQUENCE_IDS:
107        raise ValueError(f"'{sequence_id}' is not a valid sequence id. Choose one of {SEQUENCE_IDS}.")
108
109    sequence_dirs = natsorted(glob(os.path.join(path, f"{sequence_id}-AR-*")))
110    if len(sequence_dirs) == len(ANISOTROPY_RATIOS):
111        return sequence_dirs
112
113    os.makedirs(path, exist_ok=True)
114
115    zip_path = os.path.join(path, f"{sequence_id}.zip")
116    util.download_source(path=zip_path, url=URLS[sequence_id], download=download, checksum=CHECKSUMS[sequence_id])
117    _unzip_with_offset_fix(zip_path, path)
118
119    # Some archives (e.g. 'OE-ID350') contain one nested archive per anisotropy ratio, which we extract as well.
120    for nested_zip_path in natsorted(glob(os.path.join(path, f"{sequence_id}-AR-*.zip"))):
121        util.unzip(zip_path=nested_zip_path, dst=path)
122
123    # The archives contain macOS metadata folders which we do not need.
124    macos_dir = os.path.join(path, "__MACOSX")
125    if os.path.exists(macos_dir):
126        shutil.rmtree(macos_dir)
127
128    sequence_dirs = natsorted(glob(os.path.join(path, f"{sequence_id}-AR-*")))
129    assert len(sequence_dirs) == len(ANISOTROPY_RATIOS), f"Unexpected folder structure for '{sequence_id}'."
130
131    return sequence_dirs
132
133
134def get_bbbc046_paths(
135    path: Union[os.PathLike, str],
136    sequence_ids: Optional[Sequence[str]] = None,
137    anisotropy_ratio: Optional[Literal[1, 2, 4, 8]] = None,
138    fluorescence_factor: Optional[Literal["0.25", "0.50", "1.00", "2.00", "4.00"]] = None,
139    download: bool = False,
140) -> Tuple[List[str], List[str]]:
141    """Get paths to the BBBC046 data.
142
143    Args:
144        path: Filepath to a folder where the data is downloaded for further processing.
145        sequence_ids: The base sequences to use. By default, all 9 base sequences are used.
146        anisotropy_ratio: The anisotropy ratio to use. By default, all 4 anisotropy ratios are used.
147        fluorescence_factor: The fluorescence level factor to use. By default, all 5 factors are used.
148        download: Whether to download the data if it is not present.
149
150    Returns:
151        List of filepaths for the image data.
152        List of filepaths for the label data.
153    """
154    if sequence_ids is None:
155        sequence_ids = SEQUENCE_IDS
156    elif isinstance(sequence_ids, str):
157        sequence_ids = [sequence_ids]
158
159    if anisotropy_ratio is not None and anisotropy_ratio not in ANISOTROPY_RATIOS:
160        raise ValueError(f"'{anisotropy_ratio}' is not a valid anisotropy ratio. Choose one of {ANISOTROPY_RATIOS}.")
161    if fluorescence_factor is not None and fluorescence_factor not in FLUORESCENCE_FACTORS:
162        raise ValueError(
163            f"'{fluorescence_factor}' is not a valid fluorescence factor. Choose one of {FLUORESCENCE_FACTORS}."
164        )
165
166    ar_pattern = "*" if anisotropy_ratio is None else str(anisotropy_ratio)
167    factor_pattern = "*" if fluorescence_factor is None else fluorescence_factor
168
169    raw_paths, label_paths = [], []
170    for sequence_id in sequence_ids:
171        get_bbbc046_data(path, sequence_id, download)
172        pattern = os.path.join(path, f"{sequence_id}-AR-{ar_pattern}", f"factor-{factor_pattern}", "img_t*.tif")
173        curr_raw_paths = natsorted(glob(pattern))
174        assert len(curr_raw_paths) > 0, f"No volumes found for '{sequence_id}' with the pattern '{pattern}'."
175
176        for raw_path in curr_raw_paths:
177            label_path = os.path.join(
178                os.path.dirname(os.path.dirname(raw_path)), os.path.basename(raw_path).replace("img_", "mask_")
179            )
180            assert os.path.exists(label_path), f"The label volume '{label_path}' is missing."
181            raw_paths.append(raw_path)
182            label_paths.append(label_path)
183
184    return raw_paths, label_paths
185
186
187def get_bbbc046_dataset(
188    path: Union[os.PathLike, str],
189    patch_shape: Tuple[int, ...],
190    sequence_ids: Optional[Sequence[str]] = None,
191    anisotropy_ratio: Optional[Literal[1, 2, 4, 8]] = None,
192    fluorescence_factor: Optional[Literal["0.25", "0.50", "1.00", "2.00", "4.00"]] = None,
193    resize_inputs: bool = False,
194    download: bool = False,
195    **kwargs
196) -> Dataset:
197    """Get the BBBC046 dataset for cell body and filopodia segmentation in synthetic 3D time-lapse images.
198
199    Args:
200        path: Filepath to a folder where the data is downloaded for further processing.
201        patch_shape: The patch shape to use for training.
202        sequence_ids: The base sequences to use. By default, all 9 base sequences are used.
203        anisotropy_ratio: The anisotropy ratio to use. By default, all 4 anisotropy ratios are used.
204        fluorescence_factor: The fluorescence level factor to use. By default, all 5 factors are used.
205        resize_inputs: Whether to resize inputs to the desired patch shape.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
208
209    Returns:
210        The segmentation dataset.
211    """
212    raw_paths, label_paths = get_bbbc046_paths(path, sequence_ids, anisotropy_ratio, fluorescence_factor, download)
213
214    if resize_inputs:
215        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
216        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
217            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
218        )
219
220    return torch_em.default_segmentation_dataset(
221        raw_paths=raw_paths,
222        raw_key=None,
223        label_paths=label_paths,
224        label_key=None,
225        patch_shape=patch_shape,
226        is_seg_dataset=True,
227        **kwargs
228    )
229
230
231def get_bbbc046_loader(
232    path: Union[os.PathLike, str],
233    batch_size: int,
234    patch_shape: Tuple[int, ...],
235    sequence_ids: Optional[Sequence[str]] = None,
236    anisotropy_ratio: Optional[Literal[1, 2, 4, 8]] = None,
237    fluorescence_factor: Optional[Literal["0.25", "0.50", "1.00", "2.00", "4.00"]] = None,
238    resize_inputs: bool = False,
239    download: bool = False,
240    **kwargs
241) -> DataLoader:
242    """Get the BBBC046 dataloader for cell body and filopodia segmentation in synthetic 3D time-lapse images.
243
244    Args:
245        path: Filepath to a folder where the data is downloaded for further processing.
246        batch_size: The batch size for training.
247        patch_shape: The patch shape to use for training.
248        sequence_ids: The base sequences to use. By default, all 9 base sequences are used.
249        anisotropy_ratio: The anisotropy ratio to use. By default, all 4 anisotropy ratios are used.
250        fluorescence_factor: The fluorescence level factor to use. By default, all 5 factors are used.
251        resize_inputs: Whether to resize inputs to the desired patch shape.
252        download: Whether to download the data if it is not present.
253        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
254
255    Returns:
256        The DataLoader.
257    """
258    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
259    dataset = get_bbbc046_dataset(
260        path, patch_shape, sequence_ids, anisotropy_ratio, fluorescence_factor, resize_inputs, download, **ds_kwargs
261    )
262    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)