torch_em.data.datasets.util

  1import os
  2import hashlib
  3import inspect
  4import zipfile
  5import requests
  6from tqdm import tqdm
  7from warnings import warn
  8from subprocess import run
  9from xml.dom import minidom
 10from packaging import version
 11from shutil import copyfileobj, which
 12
 13from typing import Optional, Tuple, Literal
 14
 15import numpy as np
 16from skimage.draw import polygon
 17
 18import torch
 19
 20import torch_em
 21from torch_em.transform import get_raw_transform
 22from torch_em.transform.generic import ResizeLongestSideInputs, Compose
 23
 24try:
 25    import gdown
 26except ImportError:
 27    gdown = None
 28
 29try:
 30    from tcia_utils import nbia
 31except ModuleNotFoundError:
 32    nbia = None
 33
 34try:
 35    from cryoet_data_portal import Client, Dataset
 36except ImportError:
 37    Client, Dataset = None, None
 38
 39try:
 40    import synapseclient
 41    import synapseutils
 42except ImportError:
 43    synapseclient, synapseutils = None, None
 44
 45
 46BIOIMAGEIO_IDS = {
 47    "covid_if": "ilastik/covid_if_training_data",
 48    "cremi": "ilastik/cremi_training_data",
 49    "dsb": "ilastik/stardist_dsb_training_data",
 50    "hpa": "",  # not on bioimageio yet
 51    "isbi2012": "ilastik/isbi2012_neuron_segmentation_challenge",
 52    "kasthuri": "",  # not on bioimageio yet:
 53    "livecell": "ilastik/livecell_dataset",
 54    "lucchi": "",  # not on bioimageio yet:
 55    "mitoem": "ilastik/mitoem_segmentation_challenge",
 56    "monuseg": "deepimagej/monuseg_digital_pathology_miccai2018",
 57    "ovules": "",  # not on bioimageio yet
 58    "plantseg_root": "ilastik/plantseg_root",
 59    "plantseg_ovules": "ilastik/plantseg_ovules",
 60    "platynereis": "ilastik/platynereis_em_training_data",
 61    "snemi": "",  # not on bioimagegio yet
 62    "uro_cell": "",  # not on bioimageio yet: https://doi.org/10.1016/j.compbiomed.2020.103693
 63    "vnc": "ilastik/vnc",
 64}
 65"""@private
 66"""
 67
 68
 69def get_bioimageio_dataset_id(dataset_name):
 70    """@private
 71    """
 72    assert dataset_name in BIOIMAGEIO_IDS
 73    return BIOIMAGEIO_IDS[dataset_name]
 74
 75
 76def get_checksum(filename: str) -> str:
 77    """Get the SHA256 checksum of a file.
 78
 79    Args:
 80        filename: The filepath.
 81
 82    Returns:
 83        The checksum.
 84    """
 85    # The file is hashed in chunks, so that datasets with multi-GB archives do not run out of memory.
 86    hasher = hashlib.sha256()
 87    with open(filename, "rb") as f:
 88        for chunk in iter(lambda: f.read(64 * 1024 * 1024), b""):
 89            hasher.update(chunk)
 90    return hasher.hexdigest()
 91
 92
 93def _check_checksum(path, checksum):
 94    if checksum is not None:
 95        this_checksum = get_checksum(path)
 96        if this_checksum != checksum:
 97            raise RuntimeError(
 98                "The checksum of the download does not match the expected checksum."
 99                f"Expected: {checksum}, got: {this_checksum}"
100            )
101        print("Download successful and checksums agree.")
102    else:
103        warn("The file was downloaded, but no checksum was provided, so the file may be corrupted.")
104
105
106# this needs to be extended to support download from s3 via boto,
107# if we get a resource that is available via s3 without support for http
108def download_source(path: str, url: str, download: bool, checksum: Optional[str] = None, verify: bool = True) -> None:
109    """Download data via https.
110
111    Args:
112        path: The path for saving the data.
113        url: The url of the data.
114        download: Whether to download the data if it is not saved at `path` yet.
115        checksum: The expected checksum of the data.
116        verify: Whether to verify the https address.
117    """
118    if os.path.exists(path):
119        return
120    if not download:
121        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
122
123    # The data is downloaded to a temporary path and only moved to `path` once it is complete and verified.
124    # Otherwise an interrupted download would be mistaken for a complete one by the check above.
125    tmp_path = f"{path}.incomplete"
126    with requests.get(url, stream=True, allow_redirects=True, verify=verify) as r:
127        r.raise_for_status()  # check for error
128        # Compute checksums on the file content rather than its HTTP transfer encoding.
129        r.raw.decode_content = True
130        file_size = int(r.headers.get("Content-Length", 0))
131        desc = f"Download {url} to {path}"
132        if file_size == 0:
133            desc += " (unknown file size)"
134        with tqdm.wrapattr(r.raw, "read", total=file_size, desc=desc) as r_raw, open(tmp_path, "wb") as f:
135            copyfileobj(r_raw, f)
136
137    _check_checksum(tmp_path, checksum)
138    os.replace(tmp_path, path)
139
140
141def download_source_gdrive(
142    path: str,
143    url: str,
144    download: bool,
145    checksum: Optional[str] = None,
146    download_type: Literal["zip", "folder"] = "zip",
147    expected_samples: int = 10000,
148    quiet: bool = True,
149) -> None:
150    """Download data from google drive.
151
152    Args:
153        path: The path for saving the data.
154        url: The url of the data.
155        download: Whether to download the data if it is not saved at `path` yet.
156        checksum: The expected checksum of the data.
157        download_type: The download type, either 'zip' or 'folder'.
158        expected_samples: The maximal number of samples in the folder.
159        quiet: Whether to download quietly.
160    """
161    if os.path.exists(path):
162        return
163
164    if not download:
165        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
166
167    if gdown is None:
168        raise RuntimeError(
169            "Need gdown library to download data from google drive. "
170            "Please install gdown: 'conda install -c conda-forge gdown==4.6.3'."
171        )
172
173    print("Downloading the files. Might take a few minutes...")
174
175    if download_type == "zip":
176        gdown.download(url, path, quiet=quiet)
177        _check_checksum(path, checksum)
178    elif download_type == "folder":
179        assert version.parse(gdown.__version__) == version.parse("4.6.3"), "Please install 'gdown==4.6.3'."
180        gdown.download_folder.__globals__["MAX_NUMBER_FILES"] = expected_samples
181        gdown.download_folder(url=url, output=path, quiet=quiet, remaining_ok=True)
182    else:
183        raise ValueError("`download_path` argument expects either `zip`/`folder`")
184
185    print("Download completed.")
186
187
188def download_source_empiar(path: str, access_id: str, download: bool) -> str:
189    """Download data from EMPIAR.
190
191    Requires the ascp command from the aspera CLI.
192
193    Args:
194        path: The path for saving the data.
195        access_id: The EMPIAR accession id of the data to download.
196        download: Whether to download the data if it is not saved at `path` yet.
197
198    Returns:
199        The path to the downloaded data.
200    """
201    download_path = os.path.join(path, access_id)
202
203    if os.path.exists(download_path):
204        return download_path
205    if not download:
206        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
207
208    if which("ascp") is None:
209        raise RuntimeError(
210            "Need aspera-cli to download data from empiar. You can install it via 'conda install -c hcc aspera-cli'."
211        )
212
213    key_file = os.path.expanduser("~/.aspera/cli/etc/asperaweb_id_dsa.openssh")
214    if not os.path.exists(key_file):
215        conda_root = os.environ["CONDA_PREFIX"]
216        key_file = os.path.join(conda_root, "etc/asperaweb_id_dsa.openssh")
217
218    if not os.path.exists(key_file):
219        raise RuntimeError("Could not find the aspera ssh keyfile")
220
221    cmd = ["ascp", "-QT", "-l", "200M", "-P33001", "-i", key_file, f"emp_ext2@fasp.ebi.ac.uk:/{access_id}", path]
222    run(cmd)
223
224    return download_path
225
226
227def download_source_kaggle(path: str, dataset_name: str, download: bool, competition: bool = False):
228    """Download data from Kaggle.
229
230    Requires the Kaggle API.
231
232    Args:
233        path: The path for saving the data.
234        dataset_name: The name of the dataset to download.
235        download: Whether to download the data if it is not saved at `path` yet.
236        competition: Whether this data is from a competition and requires the kaggle.competition API.
237    """
238    if not download:
239        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
240
241    try:
242        from kaggle.api.kaggle_api_extended import KaggleApi
243    except ModuleNotFoundError:
244        msg = "Please install the Kaggle API. You can do this using 'pip install kaggle'. "
245        msg += "After you have installed kaggle, you would need an API token. "
246        msg += "Follow the instructions at https://www.kaggle.com/docs/api."
247        raise ModuleNotFoundError(msg)
248
249    api = KaggleApi()
250    api.authenticate()
251
252    if competition:
253        api.competition_download_files(competition=dataset_name, path=path, quiet=False)
254    else:
255        api.dataset_download_files(dataset=dataset_name, path=path, quiet=False)
256
257
258def download_source_tcia(path, url, dst, csv_filename, download):
259    """Download data from TCIA.
260
261    Requires the tcia_utils python package.
262
263    Args:
264        path: The path for saving the data.
265        url: The URL to the TCIA dataset.
266        dst:
267        csv_filename:
268        download: Whether to download the data if it is not saved at `path` yet.
269    """
270    if nbia is None:
271        raise RuntimeError("Requires the tcia_utils python package.")
272    if not download:
273        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
274    assert url.endswith(".tcia"), f"{url} is not a TCIA Manifest."
275
276    # Downloads the manifest file from the collection page.
277    manifest = requests.get(url=url)
278    with open(path, "wb") as f:
279        f.write(manifest.content)
280
281    # This part extracts the UIDs from the manifests and downloads them.
282    nbia.downloadSeries(series_data=path, input_type="manifest", path=dst, csv_filename=csv_filename)
283
284
285def download_source_synapse(path: str, entity: str, download: bool) -> None:
286    """Download data from synapse.
287
288    Requires the synapseclient python library.
289
290    Args:
291        path: The path for saving the data.
292        entity: The name of the data to download from synapse.
293        download: Whether to download the data if it is not saved at `path` yet.
294    """
295    if not download:
296        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
297
298    if synapseclient is None:
299        raise RuntimeError(
300            "You must install 'synapseclient' to download files from 'synapse'. "
301            "Remember to create an account and generate an authentication code for your account. "
302            "Please follow the documentation for details on creating the '~/.synapseConfig' file here: "
303            "https://python-docs.synapse.org/tutorials/authentication/."
304        )
305
306    assert entity.startswith("syn"), "The entity name does not look as expected. It should be something like 'syn123'."
307
308    # Download all files in the folder.
309    syn = synapseclient.Synapse()
310    syn.login()  # Since we do not pass any credentials here, it fetches all details from '~/.synapseConfig'.
311    synapseutils.syncFromSynapse(syn=syn, entity=entity, path=path)
312
313
314def update_kwargs(kwargs, key, value, msg=None):
315    """@private
316    """
317    if key in kwargs:
318        msg = f"{key} will be over-ridden in loader kwargs." if msg is None else msg
319        warn(msg)
320    kwargs[key] = value
321    return kwargs
322
323
324def unzip_tarfile(tar_path: str, dst: str, remove: bool = True) -> None:
325    """Unpack a tar archive.
326
327    Args:
328        tar_path: Path to the tar file.
329        dst: Where to unpack the archive.
330        remove: Whether to remove the tar file after unpacking.
331    """
332    import tarfile
333
334    if tar_path.endswith(".tar.gz") or tar_path.endswith(".tgz"):
335        access_mode = "r:gz"
336    elif tar_path.endswith(".tar"):
337        access_mode = "r:"
338    else:
339        raise ValueError(f"The provided file isn't a supported archive to unpack. Please check the file: {tar_path}.")
340
341    tar = tarfile.open(tar_path, access_mode)
342    tar.extractall(dst)
343    tar.close()
344
345    if remove:
346        os.remove(tar_path)
347
348
349def unzip_rarfile(rar_path: str, dst: str, remove: bool = True, use_rarfile: bool = True) -> None:
350    """Unpack a rar archive.
351
352    Args:
353        rar_path: Path to the rar file.
354        dst: Where to unpack the archive.
355        remove: Whether to remove the tar file after unpacking.
356        use_rarfile: Whether to use the rarfile library or aspose.zip.
357    """
358    def _extract_with_rarfile():
359        import rarfile
360        with rarfile.RarFile(rar_path) as archive:
361            archive.extractall(path=dst)
362
363    def _extract_with_aspose():
364        import aspose.zip as az
365        with az.rar.RarArchive(rar_path) as archive:
366            archive.extract_to_directory(dst)
367
368    extractors = [
369        ('rarfile', _extract_with_rarfile), ('aspose.zip', _extract_with_aspose),
370    ] if use_rarfile else [('aspose.zip', _extract_with_aspose)]
371
372    errors = []
373    for name, extractor in extractors:
374        try:
375            extractor()
376            break
377        except Exception as err:
378            errors.append((name, err))
379            if len(errors) < len(extractors):
380                next_name = extractors[len(errors)][0]
381                warn(f"Extraction with '{name}' failed for {rar_path} ({err}). Falling back to '{next_name}'.")
382    else:
383        backends = ', '.join(f"'{name}'" for name, _ in extractors)
384        raise RuntimeError(
385            f"Failed to extract rar archive {rar_path} with {backends}. "
386            "Please ensure one of the supported backends is installed and can read this archive."
387        ) from errors[-1][1]
388
389    if remove:
390        os.remove(rar_path)
391
392
393def unzip(zip_path: str, dst: str, remove: bool = True) -> None:
394    """Unpack a zip archive.
395
396    Args:
397        zip_path: Path to the zip file.
398        dst: Where to unpack the archive.
399        remove: Whether to remove the tar file after unpacking.
400    """
401    with zipfile.ZipFile(zip_path, "r") as f:
402        f.extractall(dst)
403    if remove:
404        os.remove(zip_path)
405
406
407def split_kwargs(function, **kwargs):
408    """@private
409    """
410    function_parameters = inspect.signature(function).parameters
411    parameter_names = list(function_parameters.keys())
412    other_kwargs = {k: v for k, v in kwargs.items() if k not in parameter_names}
413    kwargs = {k: v for k, v in kwargs.items() if k in parameter_names}
414    return kwargs, other_kwargs
415
416
417# this adds the default transforms for 'raw_transform' and 'transform'
418# in case these were not specified in the kwargs
419# this is NOT necessary if 'default_segmentation_dataset' is used, only if a dataset class
420# is used directly, e.g. in the LiveCell Loader
421def ensure_transforms(ndim, **kwargs):
422    """@private
423    """
424    if "raw_transform" not in kwargs:
425        kwargs = update_kwargs(kwargs, "raw_transform", torch_em.transform.get_raw_transform())
426    if "transform" not in kwargs:
427        kwargs = update_kwargs(kwargs, "transform", torch_em.transform.get_augmentations(ndim=ndim))
428    return kwargs
429
430
431def add_instance_label_transform(
432    kwargs, add_binary_target, label_dtype=None, binary=False, boundaries=False, offsets=None, binary_is_exclusive=True,
433):
434    """@private
435    """
436    if binary_is_exclusive:
437        assert sum((offsets is not None, boundaries, binary)) <= 1
438    else:
439        assert sum((offsets is not None, boundaries)) <= 1
440    if offsets is not None:
441        label_transform2 = torch_em.transform.label.AffinityTransform(offsets=offsets,
442                                                                      add_binary_target=add_binary_target,
443                                                                      add_mask=True)
444        msg = "Offsets are passed, but 'label_transform2' is in the kwargs. It will be over-ridden."
445        kwargs = update_kwargs(kwargs, "label_transform2", label_transform2, msg=msg)
446        label_dtype = torch.float32
447    elif boundaries:
448        label_transform = torch_em.transform.label.BoundaryTransform(add_binary_target=add_binary_target)
449        msg = "Boundaries is set to true, but 'label_transform' is in the kwargs. It will be over-ridden."
450        kwargs = update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
451        label_dtype = torch.float32
452    elif binary:
453        label_transform = torch_em.transform.label.labels_to_binary
454        msg = "Binary is set to true, but 'label_transform' is in the kwargs. It will be over-ridden."
455        kwargs = update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
456        label_dtype = torch.float32
457    return kwargs, label_dtype
458
459
460def update_kwargs_for_resize_trafo(kwargs, patch_shape, resize_inputs, resize_kwargs=None, ensure_rgb=None):
461    """@private
462    """
463    # Checks for raw_transform and label_transform incoming values.
464    # If yes, it will automatically merge these two transforms to apply them together.
465    if resize_inputs:
466        assert isinstance(resize_kwargs, dict)
467
468        target_shape = resize_kwargs.get("patch_shape")
469        if len(resize_kwargs["patch_shape"]) == 3:
470            # we only need the XY dimensions to reshape the inputs along them.
471            target_shape = target_shape[1:]
472            # we provide the Z dimension value to return the desired number of slices and not the whole volume
473            kwargs["z_ext"] = resize_kwargs["patch_shape"][0]
474
475        raw_trafo = ResizeLongestSideInputs(target_shape=target_shape, is_rgb=resize_kwargs["is_rgb"])
476        label_trafo = ResizeLongestSideInputs(target_shape=target_shape, is_label=True)
477
478        # The patch shape provided to the dataset. Here, "None" means that the entire volume will be loaded.
479        patch_shape = None
480
481    if ensure_rgb is None:
482        raw_trafos = []
483    else:
484        assert not isinstance(ensure_rgb, bool), "'ensure_rgb' is expected to be a function."
485        raw_trafos = [ensure_rgb]
486
487    if "raw_transform" in kwargs:
488        raw_trafos.extend([raw_trafo, kwargs["raw_transform"]])
489    else:
490        raw_trafos.extend([raw_trafo, get_raw_transform()])
491
492    kwargs["raw_transform"] = Compose(*raw_trafos, is_multi_tensor=False)
493
494    if "label_transform" in kwargs:
495        trafo = Compose(label_trafo, kwargs["label_transform"], is_multi_tensor=False)
496        kwargs["label_transform"] = trafo
497    else:
498        kwargs["label_transform"] = label_trafo
499
500    return kwargs, patch_shape
501
502
503def generate_labeled_array_from_xml(shape: Tuple[int, ...], xml_file: str) -> np.ndarray:
504    """Generate a label mask from a contour defined in a xml annotation file.
505
506    Function taken from: https://github.com/rshwndsz/hover-net/blob/master/lightning_hovernet.ipynb
507
508    Args:
509        shape: The image shape.
510        xml_file: The path to the xml file with contour annotations.
511
512    Returns:
513        The label mask.
514    """
515    # DOM object created by the minidom parser
516    xDoc = minidom.parse(xml_file)
517
518    # List of all Region tags
519    regions = xDoc.getElementsByTagName('Region')
520
521    # List which will store the vertices for each region
522    xy = []
523    for region in regions:
524        # Loading all the vertices in the region
525        vertices = region.getElementsByTagName('Vertex')
526
527        # The vertices of a region will be stored in a array
528        vw = np.zeros((len(vertices), 2))
529
530        for index, vertex in enumerate(vertices):
531            # Storing the values of x and y coordinate after conversion
532            vw[index][0] = float(vertex.getAttribute('X'))
533            vw[index][1] = float(vertex.getAttribute('Y'))
534
535        # Append the vertices of a region
536        xy.append(np.int32(vw))
537
538    # Creating a completely black image
539    mask = np.zeros(shape, np.float32)
540
541    for i, contour in enumerate(xy):
542        r, c = polygon(np.array(contour)[:, 1], np.array(contour)[:, 0], shape=shape)
543        mask[r, c] = i
544    return mask
545
546
547# This function could be extended to convert WSIs (or modalities with multiple resolutions).
548def convert_svs_to_array(
549    path: str, location: Tuple[int, int] = (0, 0), level: int = 0, img_size: Tuple[int, int] = None,
550) -> np.ndarray:
551    """Convert a .svs file for WSI imagging to a numpy array.
552
553    Requires the tiffslide python library.
554    The function can load multi-resolution images. You can specify the resolution level via `level`.
555
556    Args:
557        path: File path ath to the svs file.
558        location: Pixel location (x, y) in level 0 of the image.
559        level: Target level used to read the image.
560        img_size: Size of the image. If None, the shape of the image at `level` is used.
561
562    Returns:
563        The image as numpy array.
564    """
565    from tiffslide import TiffSlide
566
567    assert path.endswith(".svs"), f"The provided file ({path}) isn't in svs format"
568    _slide = TiffSlide(path)
569    if img_size is None:
570        img_size = _slide.level_dimensions[0]
571    return _slide.read_region(location=location, level=level, size=img_size, as_array=True)
572
573
574def download_from_cryo_et_portal(path: str, dataset_id: int, download: bool) -> str:
575    """Download data from the CryoET Data Portal.
576
577    Requires the cryoet-data-portal python library.
578
579    Args:
580        path: The path for saving the data.
581        dataset_id: The id of the data to download from the portal.
582        download: Whether to download the data if it is not saved at `path` yet.
583
584    Returns:
585        The file path to the downloaded data.
586    """
587    if Client is None or Dataset is None:
588        raise RuntimeError("Please install CryoETDataPortal via 'pip install cryoet-data-portal'")
589
590    output_path = os.path.join(path, str(dataset_id))
591    if os.path.exists(output_path):
592        return output_path
593
594    if not download:
595        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
596
597    client = Client()
598    dataset = Dataset.get_by_id(client, dataset_id)
599    dataset.download_everything(dest_path=path)
600
601    return output_path