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 unzip_7z(path_7z: str, dst: str, remove: bool = True) -> None:
408    """Unpack a 7z archive.
409
410    Args:
411        path_7z: Path to the 7z file.
412        dst: Where to unpack the archive.
413        remove: Whether to remove the 7z file after unpacking.
414    """
415    if which("7z") is None:
416        raise RuntimeError("Need the 'p7zip' CLI to extract 7z archives. You can install it via 'conda install -c conda-forge p7zip'.")  # noqa
417
418    run(["7z", "x", f"-o{dst}", "-y", path_7z])
419
420    if remove:
421        os.remove(path_7z)
422
423
424def split_kwargs(function, **kwargs):
425    """@private
426    """
427    function_parameters = inspect.signature(function).parameters
428    parameter_names = list(function_parameters.keys())
429    other_kwargs = {k: v for k, v in kwargs.items() if k not in parameter_names}
430    kwargs = {k: v for k, v in kwargs.items() if k in parameter_names}
431    return kwargs, other_kwargs
432
433
434# this adds the default transforms for 'raw_transform' and 'transform'
435# in case these were not specified in the kwargs
436# this is NOT necessary if 'default_segmentation_dataset' is used, only if a dataset class
437# is used directly, e.g. in the LiveCell Loader
438def ensure_transforms(ndim, **kwargs):
439    """@private
440    """
441    if "raw_transform" not in kwargs:
442        kwargs = update_kwargs(kwargs, "raw_transform", torch_em.transform.get_raw_transform())
443    if "transform" not in kwargs:
444        kwargs = update_kwargs(kwargs, "transform", torch_em.transform.get_augmentations(ndim=ndim))
445    return kwargs
446
447
448def add_instance_label_transform(
449    kwargs, add_binary_target, label_dtype=None, binary=False, boundaries=False, offsets=None, binary_is_exclusive=True,
450):
451    """@private
452    """
453    if binary_is_exclusive:
454        assert sum((offsets is not None, boundaries, binary)) <= 1
455    else:
456        assert sum((offsets is not None, boundaries)) <= 1
457    if offsets is not None:
458        label_transform2 = torch_em.transform.label.AffinityTransform(offsets=offsets,
459                                                                      add_binary_target=add_binary_target,
460                                                                      add_mask=True)
461        msg = "Offsets are passed, but 'label_transform2' is in the kwargs. It will be over-ridden."
462        kwargs = update_kwargs(kwargs, "label_transform2", label_transform2, msg=msg)
463        label_dtype = torch.float32
464    elif boundaries:
465        label_transform = torch_em.transform.label.BoundaryTransform(add_binary_target=add_binary_target)
466        msg = "Boundaries is set to true, but 'label_transform' is in the kwargs. It will be over-ridden."
467        kwargs = update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
468        label_dtype = torch.float32
469    elif binary:
470        label_transform = torch_em.transform.label.labels_to_binary
471        msg = "Binary is set to true, but 'label_transform' is in the kwargs. It will be over-ridden."
472        kwargs = update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
473        label_dtype = torch.float32
474    return kwargs, label_dtype
475
476
477def update_kwargs_for_resize_trafo(kwargs, patch_shape, resize_inputs, resize_kwargs=None, ensure_rgb=None):
478    """@private
479    """
480    # Checks for raw_transform and label_transform incoming values.
481    # If yes, it will automatically merge these two transforms to apply them together.
482    if resize_inputs:
483        assert isinstance(resize_kwargs, dict)
484
485        target_shape = resize_kwargs.get("patch_shape")
486        if len(resize_kwargs["patch_shape"]) == 3:
487            # we only need the XY dimensions to reshape the inputs along them.
488            target_shape = target_shape[1:]
489            # we provide the Z dimension value to return the desired number of slices and not the whole volume
490            kwargs["z_ext"] = resize_kwargs["patch_shape"][0]
491
492        raw_trafo = ResizeLongestSideInputs(target_shape=target_shape, is_rgb=resize_kwargs["is_rgb"])
493        label_trafo = ResizeLongestSideInputs(target_shape=target_shape, is_label=True)
494
495        # The patch shape provided to the dataset. Here, "None" means that the entire volume will be loaded.
496        patch_shape = None
497
498    if ensure_rgb is None:
499        raw_trafos = []
500    else:
501        assert not isinstance(ensure_rgb, bool), "'ensure_rgb' is expected to be a function."
502        raw_trafos = [ensure_rgb]
503
504    if "raw_transform" in kwargs:
505        raw_trafos.extend([raw_trafo, kwargs["raw_transform"]])
506    else:
507        raw_trafos.extend([raw_trafo, get_raw_transform()])
508
509    kwargs["raw_transform"] = Compose(*raw_trafos, is_multi_tensor=False)
510
511    if "label_transform" in kwargs:
512        trafo = Compose(label_trafo, kwargs["label_transform"], is_multi_tensor=False)
513        kwargs["label_transform"] = trafo
514    else:
515        kwargs["label_transform"] = label_trafo
516
517    return kwargs, patch_shape
518
519
520def generate_labeled_array_from_xml(shape: Tuple[int, ...], xml_file: str) -> np.ndarray:
521    """Generate a label mask from a contour defined in a xml annotation file.
522
523    Function taken from: https://github.com/rshwndsz/hover-net/blob/master/lightning_hovernet.ipynb
524
525    Args:
526        shape: The image shape.
527        xml_file: The path to the xml file with contour annotations.
528
529    Returns:
530        The label mask.
531    """
532    # DOM object created by the minidom parser
533    xDoc = minidom.parse(xml_file)
534
535    # List of all Region tags
536    regions = xDoc.getElementsByTagName('Region')
537
538    # List which will store the vertices for each region
539    xy = []
540    for region in regions:
541        # Loading all the vertices in the region
542        vertices = region.getElementsByTagName('Vertex')
543
544        # The vertices of a region will be stored in a array
545        vw = np.zeros((len(vertices), 2))
546
547        for index, vertex in enumerate(vertices):
548            # Storing the values of x and y coordinate after conversion
549            vw[index][0] = float(vertex.getAttribute('X'))
550            vw[index][1] = float(vertex.getAttribute('Y'))
551
552        # Append the vertices of a region
553        xy.append(np.int32(vw))
554
555    # Creating a completely black image
556    mask = np.zeros(shape, np.uint32)  # Integer instance ids; float labels break connected-component ops.
557
558    # Start the instance ids at 1: id 0 is background, so enumerating from 0 silently drops the first region.
559    for i, contour in enumerate(xy, start=1):
560        r, c = polygon(np.array(contour)[:, 1], np.array(contour)[:, 0], shape=shape)
561        mask[r, c] = i
562    return mask
563
564
565# This function could be extended to convert WSIs (or modalities with multiple resolutions).
566def convert_svs_to_array(
567    path: str, location: Tuple[int, int] = (0, 0), level: int = 0, img_size: Tuple[int, int] = None,
568) -> np.ndarray:
569    """Convert a .svs file for WSI imagging to a numpy array.
570
571    Requires the tiffslide python library.
572    The function can load multi-resolution images. You can specify the resolution level via `level`.
573
574    Args:
575        path: File path ath to the svs file.
576        location: Pixel location (x, y) in level 0 of the image.
577        level: Target level used to read the image.
578        img_size: Size of the image. If None, the shape of the image at `level` is used.
579
580    Returns:
581        The image as numpy array.
582    """
583    assert path.endswith(".svs"), f"The provided file ({path}) isn't in svs format"
584
585    try:
586        from tiffslide import TiffSlide
587    except ImportError:
588        # svs is a pyramidal TIFF variant, so tifffile can read it without the tiffslide dependency.
589        import tifffile
590        with tifffile.TiffFile(path) as f:
591            image = f.series[0].levels[level].asarray()
592        x, y = location
593        if img_size is not None:
594            image = image[y:y + img_size[1], x:x + img_size[0]]
595        else:
596            image = image[y:, x:]
597        return image
598
599    _slide = TiffSlide(path)
600    if img_size is None:
601        img_size = _slide.level_dimensions[0]
602    return _slide.read_region(location=location, level=level, size=img_size, as_array=True)
603
604
605def download_from_cryo_et_portal(path: str, dataset_id: int, download: bool) -> str:
606    """Download data from the CryoET Data Portal.
607
608    Requires the cryoet-data-portal python library.
609
610    Args:
611        path: The path for saving the data.
612        dataset_id: The id of the data to download from the portal.
613        download: Whether to download the data if it is not saved at `path` yet.
614
615    Returns:
616        The file path to the downloaded data.
617    """
618    if Client is None or Dataset is None:
619        raise RuntimeError("Please install CryoETDataPortal via 'pip install cryoet-data-portal'")
620
621    output_path = os.path.join(path, str(dataset_id))
622    if os.path.exists(output_path):
623        return output_path
624
625    if not download:
626        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
627
628    client = Client()
629    dataset = Dataset.get_by_id(client, dataset_id)
630    dataset.download_everything(dest_path=path)
631
632    return output_path