torch_em.data.datasets.light_microscopy.bbbc041

  1"""The BBBC041 dataset contains brightfield images of Giemsa-stained blood smears from patients with
  2P. vivax malaria, with bounding box annotations for each cell and its class.
  3
  4The dataset has 1,328 images: 1,208 for training and 120 for testing. Each cell is annotated with a box and
  5one of these classes: red blood cell, trophozoite, ring, schizont, gametocyte, leukocyte and difficult.
  6This loader draws the boxes into label images, so the instance and semantic labels have box shapes:
  7- 'instances': every box gets its own id.
  8- 'semantic': every box gets the id of its class, see CLASSES.
  9Boxes overlap in dense smears. The larger boxes are drawn first, so that a smaller box stays visible on top.
 10
 11The dataset is located at https://bbbc.broadinstitute.org/BBBC041.
 12This dataset is from the publication https://doi.org/10.1007/978-3-319-66179-7_1.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17import json
 18from glob import glob
 19from natsort import natsorted
 20from typing import Union, Tuple, Literal, List
 21
 22import numpy as np
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URL = "https://data.broadinstitute.org/bbbc/BBBC041/malaria.zip"
 33CHECKSUM = "46e47b02284fcab4d3f8cff291f08f0dbc1e9296c4ec974f3f4d989652446e81"
 34
 35# The semantic id of each class. The id 0 is the background.
 36CLASSES = {
 37    "red blood cell": 1,
 38    "trophozoite": 2,
 39    "ring": 3,
 40    "schizont": 4,
 41    "gametocyte": 5,
 42    "leukocyte": 6,
 43    "difficult": 7,
 44}
 45
 46SPLIT_FILES = {"train": "training.json", "test": "test.json"}
 47
 48
 49def _draw_boxes(objects, shape):
 50    """Draw the boxes of the objects into an instance and a semantic label image, the largest boxes first."""
 51    def area(obj):
 52        box = obj["bounding_box"]
 53        return (box["maximum"]["r"] - box["minimum"]["r"]) * (box["maximum"]["c"] - box["minimum"]["c"])
 54
 55    instances = np.zeros(shape, dtype="int64")
 56    semantic = np.zeros(shape, dtype="uint8")
 57    for instance_id, obj in enumerate(sorted(objects, key=area, reverse=True), start=1):
 58        box = obj["bounding_box"]
 59        r0, c0 = max(box["minimum"]["r"], 0), max(box["minimum"]["c"], 0)
 60        r1, c1 = min(box["maximum"]["r"], shape[0]), min(box["maximum"]["c"], shape[1])
 61        instances[r0:r1, c0:c1] = instance_id
 62        semantic[r0:r1, c0:c1] = CLASSES[obj["category"]]
 63    return instances, semantic
 64
 65
 66def _create_h5_data(data_dir, h5_dir, split):
 67    """Create one h5 file per image with the raw image and the instance and semantic labels."""
 68    import h5py
 69    from tqdm import tqdm
 70
 71    os.makedirs(h5_dir, exist_ok=True)
 72    with open(os.path.join(data_dir, SPLIT_FILES[split])) as f:
 73        annotations = json.load(f)
 74
 75    for annotation in tqdm(annotations, desc=f"Creating h5 files for '{split}'"):
 76        image_path = os.path.join(data_dir, annotation["image"]["pathname"].lstrip("/"))
 77        h5_path = os.path.join(h5_dir, os.path.splitext(os.path.basename(image_path))[0] + ".h5")
 78        if os.path.exists(h5_path):
 79            continue
 80
 81        raw = imageio.imread(image_path)
 82        if raw.ndim == 3:
 83            raw = raw[..., :3].transpose(2, 0, 1)
 84        instances, semantic = _draw_boxes(annotation["objects"], raw.shape[-2:])
 85
 86        with h5py.File(h5_path, "w") as f:
 87            f.create_dataset("raw", data=raw, compression="gzip")
 88            f.create_dataset("labels/instances", data=instances, compression="gzip")
 89            f.create_dataset("labels/semantic", data=semantic, compression="gzip")
 90
 91
 92def get_bbbc041_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 93    """Download the BBBC041 dataset.
 94
 95    Args:
 96        path: Filepath to a folder where the data is downloaded for further processing.
 97        download: Whether to download the data if it is not present.
 98
 99    Returns:
100        Filepath to the folder with the images and the annotation files.
101    """
102    data_dir = os.path.join(path, "malaria")
103    if os.path.exists(os.path.join(data_dir, SPLIT_FILES["train"])):
104        return data_dir
105
106    os.makedirs(path, exist_ok=True)
107    zip_path = os.path.join(path, "malaria.zip")
108    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
109    util.unzip(zip_path=zip_path, dst=path)
110
111    return data_dir
112
113
114def get_bbbc041_paths(
115    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
116) -> List[str]:
117    """Get paths to the BBBC041 data.
118
119    Args:
120        path: Filepath to a folder where the data is downloaded for further processing.
121        split: The choice of data split. Either 'train' or 'test'.
122        download: Whether to download the data if it is not present.
123
124    Returns:
125        List of filepaths for the h5 data.
126    """
127    if split not in SPLIT_FILES:
128        raise ValueError(f"'{split}' is not a valid split. Choose either 'train' or 'test'.")
129
130    data_dir = get_bbbc041_data(path, download)
131    h5_dir = os.path.join(path, "h5_data", split)
132    with open(os.path.join(data_dir, SPLIT_FILES[split])) as f:
133        n_images = len(json.load(f))
134    if len(glob(os.path.join(h5_dir, "*.h5"))) < n_images:
135        _create_h5_data(data_dir, h5_dir, split)
136
137    h5_paths = natsorted(glob(os.path.join(h5_dir, "*.h5")))
138    assert len(h5_paths) == n_images, f"Expected {n_images} h5 files for '{split}', found {len(h5_paths)}."
139    return h5_paths
140
141
142def get_bbbc041_dataset(
143    path: Union[os.PathLike, str],
144    patch_shape: Tuple[int, int],
145    split: Literal["train", "test"],
146    label_choice: Literal["instances", "semantic"] = "instances",
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> Dataset:
151    """Get the BBBC041 dataset for blood cell segmentation and classification in malaria blood smears.
152
153    Args:
154        path: Filepath to a folder where the data is downloaded for further processing.
155        patch_shape: The patch shape to use for training.
156        split: The choice of data split. Either 'train' or 'test'.
157        label_choice: The choice of labels. Either 'instances' (one id per box) or 'semantic' (the class of each
158            box, see CLASSES). Both are drawn from the bounding boxes.
159        resize_inputs: Whether to resize inputs to the desired patch shape.
160        download: Whether to download the data if it is not present.
161        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
162
163    Returns:
164        The segmentation dataset.
165    """
166    if label_choice not in ("instances", "semantic"):
167        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose either 'instances' or 'semantic'.")
168
169    h5_paths = get_bbbc041_paths(path, split, download)
170
171    if label_choice == "instances":
172        kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True, label_dtype=np.int64)
173
174    if resize_inputs:
175        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
176        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
177            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
178        )
179
180    return torch_em.default_segmentation_dataset(
181        raw_paths=h5_paths,
182        raw_key="raw",
183        label_paths=h5_paths,
184        label_key=f"labels/{label_choice}",
185        patch_shape=patch_shape,
186        with_channels=True,
187        ndim=2,
188        **kwargs
189    )
190
191
192def get_bbbc041_loader(
193    path: Union[os.PathLike, str],
194    batch_size: int,
195    patch_shape: Tuple[int, int],
196    split: Literal["train", "test"],
197    label_choice: Literal["instances", "semantic"] = "instances",
198    resize_inputs: bool = False,
199    download: bool = False,
200    **kwargs
201) -> DataLoader:
202    """Get the BBBC041 dataloader for blood cell segmentation and classification in malaria blood smears.
203
204    Args:
205        path: Filepath to a folder where the data is downloaded for further processing.
206        batch_size: The batch size for training.
207        patch_shape: The patch shape to use for training.
208        split: The choice of data split. Either 'train' or 'test'.
209        label_choice: The choice of labels. Either 'instances' (one id per box) or 'semantic' (the class of each
210            box, see CLASSES). Both are drawn from the bounding boxes.
211        resize_inputs: Whether to resize inputs to the desired patch shape.
212        download: Whether to download the data if it is not present.
213        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
214
215    Returns:
216        The DataLoader.
217    """
218    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
219    dataset = get_bbbc041_dataset(path, patch_shape, split, label_choice, resize_inputs, download, **ds_kwargs)
220    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)