torch_em.data.datasets.light_microscopy.mrose_bacteria

  1"""The M-ROSE bacteria dataset contains annotations for bacteria segmentation
  2in Gram-stained bright-field microscopy images of respiratory specimens.
  3
  4The images come from microbiological rapid on-site evaluation (M-ROSE) of patients with a lung
  5infection. The dataset holds 6005 crops of 640 x 640 pixels with 11824 bacteria. Every bacterium
  6has a polygon and a Gram status, so this loader can return instance labels or the two Gram classes.
  7
  8NOTE: The annotations are sparse. A clinical smear holds debris and unlabelled objects, and every
  9object without a polygon becomes background. Take this into account when you measure recall.
 10
 11NOTE: The archive also holds detection labels that separate cocci from bacilli. This loader uses
 12the segmentation folder only, and its labels carry the Gram status alone.
 13
 14NOTE: The image '000077_0_6' is in none of the three split files, so the loader skips it. The splits
 15cover 6004 of the 6005 crops.
 16
 17The dataset is located at https://doi.org/10.5281/zenodo.10526360 under the CC BY 4.0 license.
 18This dataset is from the publication https://doi.org/10.1038/s41597-024-03370-5.
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23import json
 24import zipfile
 25from glob import glob
 26from pathlib import Path
 27from natsort import natsorted
 28from typing import List, Literal, Optional, Tuple, Union
 29
 30import numpy as np
 31import imageio.v3 as imageio
 32
 33from torch.utils.data import DataLoader, Dataset
 34
 35import torch_em
 36
 37from .. import util
 38
 39
 40URL = "https://zenodo.org/records/10526360/files/DeepDataSet.zip?download=1"
 41CHECKSUM = "c5ecaa65fa8c515b3495148c55f48aceeecea0114b4ecbf59c6faf43fb2df4a6"
 42
 43# The archive also holds a detection folder and a plain image folder, which this loader ignores.
 44ARCHIVE_FOLDER = "SegmentationDataSet"
 45
 46SPLITS = ("train", "val", "test")
 47
 48# The annotations mark a Gram-positive bacterium with 'G+' and a Gram-negative one with 'G'.
 49GRAM_IDS = {"G+": 1, "G": 2}
 50
 51
 52def _extract_segmentation_folder(zip_path: str, path: str) -> None:
 53    """Extract only the segmentation folder of the archive."""
 54    with zipfile.ZipFile(zip_path) as archive:
 55        members = [n for n in archive.namelist() if n.startswith(f"{ARCHIVE_FOLDER}/")]
 56        if not members:
 57            raise RuntimeError(f"The archive {zip_path} does not hold a '{ARCHIVE_FOLDER}' folder.")
 58        archive.extractall(path, members=members)
 59
 60
 61def _rasterize(shapes, shape: Tuple[int, int]) -> Tuple[np.ndarray, np.ndarray]:
 62    """Draw one label per bacterium, and a second image with the Gram status."""
 63    from skimage.draw import polygon as draw_polygon
 64
 65    instances = np.zeros(shape, dtype="uint16")
 66    semantic = np.zeros(shape, dtype="uint8")
 67    for instance_id, item in enumerate(shapes, start=1):
 68        points = np.array(item["points"], dtype=float)
 69        rows, columns = draw_polygon(points[:, 1], points[:, 0], shape=shape)
 70        instances[rows, columns] = instance_id
 71        semantic[rows, columns] = GRAM_IDS.get(item.get("label"), 0)
 72    return instances, semantic
 73
 74
 75def _create_labels(data_dir: str) -> Tuple[str, str]:
 76    """Rasterize the polygons of every crop into label images."""
 77    from tqdm import tqdm
 78
 79    instance_dir = os.path.join(data_dir, "instance_labels")
 80    semantic_dir = os.path.join(data_dir, "semantic_labels")
 81    os.makedirs(instance_dir, exist_ok=True)
 82    os.makedirs(semantic_dir, exist_ok=True)
 83
 84    json_paths = natsorted(glob(os.path.join(data_dir, "json", "*.json")))
 85    for json_path in tqdm(json_paths, desc="Preprocess the M-ROSE annotations"):
 86        stem = Path(json_path).stem
 87        instance_path = os.path.join(instance_dir, f"{stem}.tif")
 88        semantic_path = os.path.join(semantic_dir, f"{stem}.tif")
 89        if os.path.exists(instance_path) and os.path.exists(semantic_path):
 90            continue
 91
 92        with open(json_path) as f:
 93            annotation = json.load(f)
 94
 95        shape = (annotation["imageHeight"], annotation["imageWidth"])
 96        instances, semantic = _rasterize(annotation.get("shapes", []), shape)
 97        imageio.imwrite(instance_path, instances, compression="zlib")
 98        imageio.imwrite(semantic_path, semantic, compression="zlib")
 99
100    return instance_dir, semantic_dir
101
102
103def get_mrose_bacteria_data(path: Union[os.PathLike, str], download: bool = False) -> str:
104    """Download the M-ROSE bacteria dataset.
105
106    Args:
107        path: Filepath to a folder where the downloaded data will be saved.
108        download: Whether to download the data if it is not present.
109
110    Returns:
111        The filepath to the extracted segmentation data.
112    """
113    data_dir = os.path.join(path, ARCHIVE_FOLDER)
114    if os.path.exists(data_dir):
115        return data_dir
116
117    os.makedirs(path, exist_ok=True)
118    zip_path = os.path.join(path, "DeepDataSet.zip")
119    util.download_source(zip_path, URL, download, CHECKSUM)
120    _extract_segmentation_folder(zip_path, path)
121
122    return data_dir
123
124
125def get_mrose_bacteria_paths(
126    path: Union[os.PathLike, str],
127    split: Literal["train", "val", "test"] = "train",
128    label_choice: Literal["instances", "semantic"] = "instances",
129    download: bool = False,
130) -> Tuple[List[str], List[str]]:
131    """Get paths to the M-ROSE bacteria data.
132
133    Args:
134        path: Filepath to a folder where the downloaded data will be saved.
135        split: The data split. Either 'train', 'val' or 'test'.
136        label_choice: The label to use. Either 'instances' for the bacteria, or 'semantic' for the
137            Gram classes, where one is Gram-positive and two is Gram-negative.
138        download: Whether to download the data if it is not present.
139
140    Returns:
141        List of filepaths for the image data.
142        List of filepaths for the label data.
143    """
144    if split not in SPLITS:
145        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
146    if label_choice not in ("instances", "semantic"):
147        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose 'instances' or 'semantic'.")
148
149    data_dir = get_mrose_bacteria_data(path, download)
150    instance_dir, semantic_dir = _create_labels(data_dir)
151    label_dir = instance_dir if label_choice == "instances" else semantic_dir
152
153    split_path = os.path.join(data_dir, "txt", f"{split}.txt")
154    with open(split_path) as f:
155        stems = [line.strip() for line in f if line.strip()]
156
157    image_paths, label_paths = [], []
158    for stem in stems:
159        image_path = os.path.join(data_dir, "images", f"{stem}.jpg")
160        label_path = os.path.join(label_dir, f"{stem}.tif")
161        if not (os.path.exists(image_path) and os.path.exists(label_path)):
162            continue
163        image_paths.append(image_path)
164        label_paths.append(label_path)
165
166    if not image_paths:
167        raise RuntimeError(f"Could not find any M-ROSE data for the '{split}' split in {data_dir}.")
168
169    return image_paths, label_paths
170
171
172def get_mrose_bacteria_dataset(
173    path: Union[os.PathLike, str],
174    patch_shape: Tuple[int, int],
175    split: Literal["train", "val", "test"] = "train",
176    label_choice: Literal["instances", "semantic"] = "instances",
177    offsets: Optional[List[List[int]]] = None,
178    boundaries: bool = False,
179    binary: bool = False,
180    download: bool = False,
181    **kwargs,
182) -> Dataset:
183    """Get the M-ROSE bacteria dataset for bacteria segmentation.
184
185    Args:
186        path: Filepath to a folder where the downloaded data will be saved.
187        patch_shape: The 2D patch shape to use for training.
188        split: The data split. Either 'train', 'val' or 'test'.
189        label_choice: The label to use. Either 'instances' for the bacteria, or 'semantic' for the
190            Gram classes, where one is Gram-positive and two is Gram-negative.
191        offsets: Offset values for affinity computation used as target.
192        boundaries: Whether to compute boundaries as the target.
193        binary: Whether to use a binary segmentation target.
194        download: Whether to download the data if it is not present.
195        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
196
197    Returns:
198        The segmentation dataset.
199    """
200    if len(patch_shape) != 2:
201        raise ValueError(f"The M-ROSE patch shape must be two-dimensional, got {patch_shape}.")
202
203    image_paths, label_paths = get_mrose_bacteria_paths(path, split, label_choice, download)
204
205    if label_choice == "instances":
206        kwargs, _ = util.add_instance_label_transform(
207            kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
208        )
209    kwargs = util.ensure_transforms(ndim=2, **kwargs)
210
211    return torch_em.default_segmentation_dataset(
212        raw_paths=image_paths,
213        raw_key=None,
214        label_paths=label_paths,
215        label_key=None,
216        patch_shape=patch_shape,
217        is_seg_dataset=False,
218        ndim=2,
219        **kwargs,
220    )
221
222
223def get_mrose_bacteria_loader(
224    path: Union[os.PathLike, str],
225    batch_size: int,
226    patch_shape: Tuple[int, int],
227    split: Literal["train", "val", "test"] = "train",
228    label_choice: Literal["instances", "semantic"] = "instances",
229    offsets: Optional[List[List[int]]] = None,
230    boundaries: bool = False,
231    binary: bool = False,
232    download: bool = False,
233    **kwargs,
234) -> DataLoader:
235    """Get the M-ROSE bacteria dataloader for bacteria segmentation.
236
237    Args:
238        path: Filepath to a folder where the downloaded data will be saved.
239        batch_size: The batch size for training.
240        patch_shape: The 2D patch shape to use for training.
241        split: The data split. Either 'train', 'val' or 'test'.
242        label_choice: The label to use. Either 'instances' for the bacteria, or 'semantic' for the
243            Gram classes, where one is Gram-positive and two is Gram-negative.
244        offsets: Offset values for affinity computation used as target.
245        boundaries: Whether to compute boundaries as the target.
246        binary: Whether to use a binary segmentation target.
247        download: Whether to download the data if it is not present.
248        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
249
250    Returns:
251        The DataLoader.
252    """
253    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
254    dataset = get_mrose_bacteria_dataset(
255        path=path,
256        patch_shape=patch_shape,
257        split=split,
258        label_choice=label_choice,
259        offsets=offsets,
260        boundaries=boundaries,
261        binary=binary,
262        download=download,
263        **ds_kwargs,
264    )
265    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)