torch_em.data.datasets.histopathology.puma

  1"""The PUMA dataset contains annotations for nucleus and tissue segmentation
  2in melanoma H&E stained histopathology images.
  3
  4This dataset is located at https://zenodo.org/records/13859989.
  5This is part of the PUMA Grand Challenge: https://puma.grand-challenge.org/.
  6The dataset is from the publication https://doi.org/10.1093/gigascience/giaf011.
  7Please cite them if you use this dataset for your research.
  8"""
  9
 10import os
 11from glob import glob
 12from tqdm import tqdm
 13from pathlib import Path
 14from natsort import natsorted
 15from typing import Union, Literal, List, Tuple
 16
 17import json
 18import numpy as np
 19import pandas as pd
 20import imageio.v3 as imageio
 21from sklearn.model_selection import train_test_split
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = {
 31    "data": "https://zenodo.org/records/15050523/files/01_training_dataset_tif_ROIs.zip",
 32    "annotations": {
 33        "nuclei": "https://zenodo.org/records/15050523/files/01_training_dataset_geojson_nuclei.zip",
 34        "tissue": "https://zenodo.org/records/15050523/files/01_training_dataset_geojson_tissue.zip",
 35    }
 36}
 37
 38CHECKSUM = {
 39    "data": "af48b879f8ff7e74b84a7114924881606f13f108aa0f9bcc21d3593b717ee022",
 40    "annotations": {
 41        "nuclei": "eda271225900d6de0759e0281f3731a570e09f2adab58bd36425b9d2dfad91a0",
 42        "tissue": "fc2835135cc28324f52eac131327f0f12c554c0b1f334a108bf4b65e0f18c42b",
 43    }
 44}
 45
 46NUCLEI_CLASS_DICT = {
 47    "nuclei_stroma": 1,
 48    "nuclei_tumor": 2,
 49    "nuclei_plasma_cell": 3,
 50    "nuclei_histiocyte": 4,
 51    "nuclei_lymphocyte": 5,
 52    "nuclei_melanophage": 6,
 53    "nuclei_neutrophil": 7,
 54    "nuclei_endothelium": 8,
 55    "nuclei_epithelium": 9,
 56    "nuclei_apoptosis": 10,
 57}
 58
 59TISSUE_CLASS_DICT = {
 60    "tissue_stroma": 1,
 61    "tissue_tumor": 2,
 62    "tissue_epidermis": 3,
 63    "tissue_blood_vessel": 4,
 64    "tissue_necrosis": 5,
 65    "tissue_white_background": 6,
 66}
 67
 68CLASS_DICT = {
 69    "nuclei": NUCLEI_CLASS_DICT,
 70    "tissue": TISSUE_CLASS_DICT,
 71}
 72
 73
 74def _get_classes_per_image(path, annotations):
 75    "Maps each ROI image id to the set of class ids present in its annotation"
 76    ann_dir = os.path.join(path, "annotations", annotations, f"01_training_dataset_geojson_{annotations}")
 77    ann_paths = glob(os.path.join(ann_dir, "*.geojson"))
 78    assert ann_paths, f"Could not find any '{annotations}' annotations in '{ann_dir}'."
 79
 80    class_dict = CLASS_DICT[annotations]
 81    classes_per_image = {}
 82    for ann_path in ann_paths:
 83        image_id = os.path.basename(ann_path).replace(f"_{annotations}.geojson", "")
 84        with open(ann_path) as f:
 85            feature_collection = json.load(f)
 86        classes_per_image[image_id] = {
 87            class_dict[feat["properties"]["classification"]["name"]]
 88            for feat in feature_collection["features"]
 89        }
 90    return classes_per_image
 91
 92
 93def _split_covers_all_classes(split_ids, classes_per_image, all_classes):
 94    "Checks that every class is present in at least one image of every split"
 95    for ids in split_ids.values():
 96        covered = set()
 97        for image_id in ids:
 98            covered.update(classes_per_image.get(image_id, set()))
 99        if not all_classes.issubset(covered):
100            return False
101    return True
102
103
104def _make_random_split(metastatic_ids, primary_ids, random_state):
105    # Create random splits per dataset: 20% for test, then 15% of the train set for val.
106    train_ids, test_ids = train_test_split(metastatic_ids, test_size=0.2, random_state=random_state)
107    train_ids, val_ids = train_test_split(train_ids, test_size=0.15, random_state=random_state)
108    # Do same as above for 'primary' samples.
109    ptrain_ids, ptest_ids = train_test_split(primary_ids, test_size=0.2, random_state=random_state)
110    ptrain_ids, pval_ids = train_test_split(ptrain_ids, test_size=0.15, random_state=random_state)
111    train_ids = train_ids + ptrain_ids
112    val_ids = val_ids + pval_ids
113    test_ids = test_ids + ptest_ids
114    return {"train": train_ids, "val": val_ids, "test": test_ids}
115
116
117def _create_split_csv(path, annotations, split, random_state=42, max_split_attempts=200):
118    "This creates a split saved to a .csv file in the dataset directory"
119    csv_path = os.path.join(path, "puma_split.csv")
120
121    if os.path.exists(csv_path):
122        df = pd.read_csv(csv_path)
123        df[split] = df[split].apply(lambda x: json.loads(x.replace("'", '"')))  # ensures all items from column in list.
124        split_list = df.iloc[0][split]
125    else:
126        print(f"Creating a new split file at '{csv_path}'.")
127        # NOTE: The ids are sorted, as 'train_test_split' is order sensitive and 'glob' does not
128        # guarantee a stable order, i.e. the fixed 'random_state' alone does not reproduce the split.
129        metastatic_ids = natsorted(
130            os.path.basename(image).split(".")[0]
131            for image in glob(os.path.join(path, "data", "01_training_dataset_tif_ROIs", "*metastatic*"))
132        )
133        primary_ids = natsorted(
134            os.path.basename(image).split(".")[0]
135            for image in glob(os.path.join(path, "data", "01_training_dataset_tif_ROIs", "*primary*"))
136        )
137
138        # The ROIs (and hence this split) are shared across annotation levels, so a split validated
139        # for one level can still drop a rare class (e.g. 'tissue_white_background') from the other.
140        coverage_checks = []
141        for level, class_dict in CLASS_DICT.items():
142            coverage_checks.append((set(class_dict.values()), _get_classes_per_image(path, level)))
143
144        split_ids, candidate = None, None
145        for attempt in range(max_split_attempts):
146            candidate = _make_random_split(metastatic_ids, primary_ids, random_state=random_state + attempt)
147            if all(
148                _split_covers_all_classes(candidate, classes_per_image, all_classes)
149                for all_classes, classes_per_image in coverage_checks
150            ):
151                split_ids = candidate
152                break
153
154        if split_ids is None:
155            print(
156                f"Warning: could not find a split covering all classes within {max_split_attempts} attempts "
157                f"(seeds {random_state}-{random_state + max_split_attempts - 1}). Using the last candidate "
158                "anyway, please check rare classes manually."
159            )
160            split_ids = candidate
161
162        df = pd.DataFrame.from_dict([split_ids])
163        df.to_csv(csv_path, index=False)
164
165        split_list = split_ids[split]
166
167    return split_list
168
169
170def _preprocess_inputs(path, annotations, split):
171    import h5py
172    try:
173        import geopandas as gpd
174    except ModuleNotFoundError:
175        raise RuntimeError("Please install 'geopandas': 'conda install -c conda-forge geopandas'.")
176
177    try:
178        from rasterio.features import rasterize
179        from rasterio.transform import from_bounds
180    except ModuleNotFoundError:
181        raise RuntimeError("Please install 'rasterio': 'conda install -c conda-forge rasterio'.")
182
183    annotation_paths = glob(
184        os.path.join(path, "annotations", annotations, f"01_training_dataset_geojson_{annotations}", "*.geojson")
185    )
186    roi_dir = os.path.join(path, "data", "01_training_dataset_tif_ROIs")
187    preprocessed_dir = os.path.join(path, split, "preprocessed")
188    os.makedirs(preprocessed_dir, exist_ok=True)
189
190    split_list = _create_split_csv(path, annotations, split)
191    print(f"The data split '{split}' has '{len(split_list)}' samples!")
192
193    for ann_path in tqdm(annotation_paths, desc=f"Preprocessing '{annotations}'"):
194        fname = os.path.basename(ann_path).replace(f"_{annotations}.geojson", ".tif")
195        image_path = os.path.join(roi_dir, fname)
196
197        # Handle inconsistent extension for sample 103 (.tiff instead of .tif).
198        if not os.path.exists(image_path):
199            image_path = image_path + "f"  # Retrying with .tiff
200
201        if os.path.basename(image_path).split(".")[0] not in split_list:
202            continue
203
204        assert os.path.exists(image_path), image_path
205
206        volume_path = os.path.join(preprocessed_dir, Path(fname).with_suffix(".h5"))
207        gdf = gpd.read_file(ann_path)
208        minx, miny, maxx, maxy = gdf.total_bounds
209
210        width, height = 1024, 1024  # roi shape
211        transform = from_bounds(minx, miny, maxx, maxy, width, height)
212
213        # Extract class ids mapped to each class name. Depending on the geopandas/pyogrio version,
214        # this property comes back either as a JSON-encoded string or already parsed into a dict.
215        class_dict = CLASS_DICT[annotations]
216        classification = gdf["classification"].apply(lambda x: json.loads(x) if isinstance(x, str) else x)
217        class_ids = [class_dict[cls_entry["name"]] for cls_entry in classification]
218        semantic_shapes = ((geom, unique_id) for geom, unique_id in zip(gdf.geometry, class_ids))
219        semantic_mask = rasterize(
220            semantic_shapes, out_shape=(height, width), transform=transform, fill=0, dtype=np.uint8
221        )
222
223        gdf['id'] = range(1, len(gdf) + 1)
224        instance_shapes = ((geom, unique_id) for geom, unique_id in zip(gdf.geometry, gdf['id']))
225        instance_mask = rasterize(
226            instance_shapes, out_shape=(height, width), transform=transform, fill=0, dtype=np.int32
227        )
228
229        # Transform labels to match expected orientation
230        instance_mask = np.flip(instance_mask)
231        instance_mask = np.fliplr(instance_mask)
232
233        semantic_mask = np.flip(semantic_mask)
234        semantic_mask = np.fliplr(semantic_mask)
235
236        image = imageio.imread(image_path)
237        image = image[..., :-1].transpose(2, 0, 1)
238
239        with h5py.File(volume_path, "a") as f:
240            if "raw" not in f.keys():
241                f.create_dataset("raw", data=image, compression="gzip")
242
243            if f"labels/instances/{annotations}" not in f.keys():
244                f.create_dataset(f"labels/instances/{annotations}", data=instance_mask, compression="gzip")
245
246            if f"labels/semantic/{annotations}" not in f.keys():
247                f.create_dataset(f"labels/semantic/{annotations}", data=semantic_mask, compression="gzip")
248
249
250def _annotations_are_stored(data_dir, annotations):
251    import h5py
252    volume_paths = glob(os.path.join(data_dir, "preprocessed", "*.h5"))
253    if not volume_paths:
254        return
255    f = h5py.File(volume_paths[0], "r")
256    return f"labels/instances/{annotations}" in f.keys()
257
258
259def get_puma_data(
260    path: Union[os.PathLike, str],
261    split: Literal["train", "val", "test"],
262    annotations: Literal['nuclei', 'tissue'] = "nuclei",
263    download: bool = False,
264) -> str:
265    """Download the PUMA data.
266
267    Args:
268        path: Filepath to a folder where the downloaded data will be saved.
269        split: The choice of data split.
270        annotations: The choice of annotations.
271        download: Whether to download the data if it is not present.
272
273    Returns:
274        Filepath where the dataset is downloaded and stored for further preprocessing.
275    """
276    if annotations not in ["nuclei", "tissue"]:
277        raise ValueError(f"'{annotations}' is not a valid annotation for the data.")
278
279    data_dir = os.path.join(path, split)
280    if os.path.exists(data_dir) and _annotations_are_stored(data_dir, annotations):
281        return data_dir
282
283    os.makedirs(path, exist_ok=True)
284
285    if not os.path.exists(os.path.join(path, "data")):
286        # Download the data.
287        zip_path = os.path.join(path, "roi.zip")
288        util.download_source(path=zip_path, url=URL["data"], download=download, checksum=CHECKSUM["data"])
289        util.unzip(zip_path=zip_path, dst=os.path.join(path, "data"))
290
291    # Download the annotations. All levels are fetched, as the split is shared across them and is
292    # validated against each level (the geojson files are small).
293    for level in CLASS_DICT:
294        annotation_dir = os.path.join(path, "annotations", level)
295        if os.path.exists(annotation_dir):
296            continue
297
298        zip_path = os.path.join(path, "annotations.zip")
299        util.download_source(
300            path=zip_path,
301            url=URL["annotations"][level],
302            download=download,
303            checksum=CHECKSUM["annotations"][level]
304        )
305        util.unzip(zip_path=zip_path, dst=annotation_dir)
306
307    _preprocess_inputs(path, annotations, split)
308
309    return data_dir
310
311
312def get_puma_paths(
313    path: Union[os.PathLike, str],
314    split: Literal["train", "val", "test"],
315    annotations: Literal['nuclei', 'tissue'] = "nuclei",
316    download: bool = False
317) -> List[str]:
318    """Get paths to the PUMA dataset.
319
320    Args:
321        path: Filepath to a folder where the downloaded data will be saved.
322        split: The choice of data split.
323        annotations: The choice of annotations.
324        download: Whether to download the data if it is not present.
325
326    Returns:
327        List of filepaths for the input data.
328    """
329    data_dir = get_puma_data(path, split, annotations, download)
330    volume_paths = natsorted(glob(os.path.join(data_dir, "preprocessed", "*.h5")))
331    return volume_paths
332
333
334def get_puma_dataset(
335    path: Union[os.PathLike, str],
336    patch_shape: Tuple[int, int],
337    split: Literal["train", "val", "test"],
338    annotations: Literal['nuclei', 'tissue'] = "nuclei",
339    label_choice: Literal["instances", "semantic"] = "instances",
340    resize_inputs: bool = False,
341    download: bool = False,
342    **kwargs
343) -> Dataset:
344    """Get the PUMA dataset for nuclei and tissue segmentation.
345
346    Args:
347        path: Filepath to a folder where the downloaded data will be saved.
348        patch_shape: The patch shape to use for training.
349        split: The choice of data split.
350        annotations: The choice of annotations.
351        label_choice: The choice of segmentation type.
352        resize_inputs: Whether to resize the inputs.
353        download: Whether to download the data if it is not present.
354        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
355
356    Returns:
357        The segmentation dataset.
358    """
359    volume_paths = get_puma_paths(path, split, annotations, download)
360
361    if resize_inputs:
362        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
363        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
364            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
365        )
366
367    return torch_em.default_segmentation_dataset(
368        raw_paths=volume_paths,
369        raw_key="raw",
370        label_paths=volume_paths,
371        label_key=f"labels/{label_choice}/{annotations}",
372        patch_shape=patch_shape,
373        with_channels=True,
374        is_seg_dataset=True,
375        ndim=2,
376        **kwargs
377    )
378
379
380def get_puma_loader(
381    path: Union[os.PathLike, str],
382    batch_size: int,
383    patch_shape: Tuple[int, int],
384    split: Literal["train", "val", "test"],
385    annotations: Literal['nuclei', 'tissue'] = "nuclei",
386    label_choice: Literal["instances", "semantic"] = "instances",
387    resize_inputs: bool = False,
388    download: bool = False,
389    **kwargs
390) -> DataLoader:
391    """Get the PUMA dataloader for nuclei and tissue segmentation.
392
393    Args:
394        path: Filepath to a folder where the downloaded data will be saved.
395        batch_size: The batch size for training.
396        patch_shape: The patch shape to use for training.
397        split: The choice of data split.
398        annotations: The choice of annotations.
399        label_choice: The choice of segmentation type.
400        resize_inputs: Whether to resize the inputs.
401        download: Whether to download the data if it is not present.
402        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
403
404    Returns:
405        The DataLoader.
406    """
407    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
408    dataset = get_puma_dataset(
409        path, patch_shape, split, annotations, label_choice, resize_inputs, download, **ds_kwargs
410    )
411    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)