torch_em.data.datasets.light_microscopy.vgg_hela

  1"""This is a dataset for counting HeLA cells in phase-contrast microscopy.
  2
  3It is described in the publication https://www.robots.ox.ac.uk/~vgg/publications/2012/Arteta12/.
  4Please cite it if you use this dataset in your research.
  5"""
  6
  7import os
  8from glob import glob
  9from shutil import rmtree
 10from typing import Tuple, Union
 11
 12import numpy as np
 13import imageio.v3 as imageio
 14from scipy.io import loadmat
 15
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23URL = "https://www.robots.ox.ac.uk/~vgg/software/cell_detection/downloads/CellDetect_v1.0.tar.gz"
 24CHECKSUM = "09825d6a8e287ddf2c4b1ef3d2f62585ec6876e3bfcd4b9bbcd3dd300e4be282"
 25
 26
 27def get_vgg_hela_data(path: Union[os.PathLike, str], download: bool) -> str:
 28    """Download the HeLA VGG dataset.
 29
 30    Args:
 31        path: Filepath to a folder where the downloaded data will be saved.
 32        download: Whether to download the data if it is not present.
 33
 34    Returns:
 35        The filepath to the training data.
 36    """
 37    os.makedirs(path, exist_ok=True)
 38    url = URL
 39    checksum = CHECKSUM
 40
 41    train_path = os.path.join(path, "train")
 42    test_path = os.path.join(path, "test")
 43
 44    if os.path.exists(train_path) and os.path.exists(test_path):
 45        return path
 46
 47    dl_path = os.path.join(path, "cell_detect.tar.gz")
 48    util.download_source(dl_path, url, download, checksum)
 49    util.unzip_tarfile(dl_path, path, True)
 50
 51    extracted_path = os.path.join(path, "CellDetect_v1.0")
 52    assert os.path.exists(extracted_path), extracted_path
 53
 54    splits_in = ["trainPhasecontrast", "testPhasecontrast"]
 55    splits_out = [train_path, test_path]
 56
 57    for split_in, out_folder in zip(splits_in, splits_out):
 58        out_im_folder = os.path.join(out_folder, "images")
 59        os.makedirs(out_im_folder, exist_ok=True)
 60
 61        out_label_folder = os.path.join(out_folder, "labels")
 62        os.makedirs(out_label_folder, exist_ok=True)
 63
 64        split_root = os.path.join(extracted_path, "phasecontrast", split_in)
 65        image_files = sorted(glob(os.path.join(split_root, "*.pgm")))
 66        mat_files = sorted(glob(os.path.join(split_root, "*.mat")))
 67
 68        for ii, (im, mat) in enumerate(zip(image_files, mat_files), 1):
 69            im = imageio.imread(im)
 70            coordinates = loadmat(mat)["gt"] - 1
 71            coordinates = (coordinates[:, 1], coordinates[:, 0])
 72
 73            out_im = os.path.join(out_im_folder, f"im{ii:02}.tif")
 74            imageio.imwrite(out_im, im, compression="zlib")
 75
 76            labels = np.zeros(im.shape, dtype="uint8")
 77            labels[coordinates] = 1
 78            out_labels = os.path.join(out_label_folder, f"im{ii:02}.tif")
 79            imageio.imwrite(out_labels, labels, compression="zlib")
 80
 81    rmtree(extracted_path)
 82    return path
 83
 84
 85def get_vgg_hela_paths(path: Union[os.PathLike, str], split: str, download: bool = False) -> Tuple[str, str]:
 86    """Get paths for HeLA VGG data.
 87
 88    Args:
 89        path: Filepath to a folder where the downloaded data will be saved.
 90        split: The split to use for the dataset. Either 'train' or 'test'.
 91        download: Whether to download the data if it is not present.
 92
 93    Returns:
 94        Filepath to the folder where image data is stored.
 95        Filepath to the folder where label data is stored.
 96    """
 97    get_vgg_hela_data(path, download)
 98
 99    image_path = os.path.join(path, split, "images")
100    label_path = os.path.join(path, split, "labels")
101
102    return image_path, label_path
103
104
105def get_vgg_hela_dataset(
106    path: Union[os.PathLike, str], split: str, patch_shape: Tuple[int, int], download: bool = False, **kwargs
107) -> Dataset:
108    """Get the HeLA VGG dataset for cell counting.
109
110    Args:
111        path: Filepath to a folder where the downloaded data will be saved.
112        split: The split to use for the dataset. Either 'train' or 'test'.
113        patch_shape: The patch shape to use for training.
114        download: Whether to download the data if it is not present.
115        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
116
117    Returns:
118       The segmentation dataset.
119    """
120    assert split in ("test", "train"), split
121
122    image_path, label_path = get_vgg_hela_paths(path, split, download)
123
124    kwargs = util.update_kwargs(kwargs, "ndim", 2)
125    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
126
127    return torch_em.default_segmentation_dataset(
128        raw_paths=image_path,
129        raw_key="*.tif",
130        label_paths=label_path,
131        label_key="*.tif",
132        patch_shape=patch_shape,
133        **kwargs
134    )
135
136
137def get_vgg_hela_loader(
138    path: Union[os.PathLike, str],
139    split: str,
140    patch_shape: Tuple[int, int],
141    batch_size: int,
142    download: bool = False,
143    **kwargs
144) -> DataLoader:
145    """Get the HeLA VGG dataloader for cell counting.
146
147    Args:
148        path: Filepath to a folder where the downloaded data will be saved.
149        split: The split to use for the dataset. Either 'train' or 'test'.
150        patch_shape: The patch shape to use for training.
151        batch_size: The batch size for training.
152        download: Whether to download the data if it is not present.
153        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
154
155    Returns:
156        The DataLoader.
157    """
158    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
159    dataset = get_vgg_hela_dataset(path, split, patch_shape, download=download, **ds_kwargs)
160    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)