torch_em.data.datasets.medical.prostatex

  1"""The PROSTATEx dataset contains annotations for prostate lesion and prostate zone segmentation
  2in multi-parametric prostate MRI.
  3
  4The MRI scans are from the SPIE-AAPM-NCI PROSTATEx challenge (346 patients) and are distributed as DICOM series
  5on TCIA. The segmentation masks are from the third-party "PROSTATEx masks" repository (Cuocolo et al. 2021):
  6- Lesion masks for 299 lesions of 200 patients, on the axial T2-weighted and on the ADC images.
  7  The lesions are stored as instance labels, the instance id is the PROSTATEx finding id.
  8- Whole gland and zonal masks on the axial T2-weighted images for 204 patients.
  9  The zone labels are: 1: peripheral zone, 2: transition zone (the rest of the gland, i.e. transition zone,
 10  central zone and anterior fibromuscular stroma).
 11
 12The label type 'zones_detailed' provides a second, finer set of zonal masks on the axial T2-weighted images
 13from the ProstateZones release (Hartman et al. 2024), which separates the zones the label type 'zones' merges
 14and adds the urethra, for 200 patients. Its label ids are 1: peripheral zone, 2: central zone,
 153: transition zone, 4: anterior fibromuscular stroma, 5: urethra (see `DETAILED_ZONE_IDS`). 40 of these
 16patients were delineated by two readers independently; the second delineation is stored next to the first as
 17'zones_detailed/t2/labels_reader2'.
 18
 19NOTE: The ProstateZones masks are distributed without the images they were drawn on and without their label
 20legend. The series was identified from the list of series that its repository ships, and the label ids were
 21assigned by measuring each annotation: the urethra is the smallest structure, the transition zone the largest,
 22the anterior fibromuscular stroma the most anterior one and the central zone the most superior one.
 23
 24NOTE: This requires the pynrrd python package to read the ProstateZones masks.
 25
 26This module downloads only the annotated T2 and ADC series from TCIA, stacks them into volumes and stores them
 27together with the aligned masks in one hdf5 file per patient. The masks were drawn on NIfTI conversions of the DICOM
 28series (dcm2niix), so the module maps them back onto the DICOM grid and checks the alignment against the NIfTI images
 29shipped with the masks. The hdf5 groups are '<label_type>/<sequence>' with the datasets 'raw' and 'labels'
 30(e.g. 'lesions/t2/raw' and 'lesions/t2/labels'); the zones group additionally contains the binary
 31whole-gland mask ('zones/t2/prostate'). Note that for some patients the lesion and the zone masks were drawn on
 32different T2 series, hence the raw data is stored per group.
 33
 34NOTE: This requires the pydicom python package.
 35
 36The dataset is located at https://www.cancerimagingarchive.net/collection/prostatex/
 37and the masks at https://github.com/rcuocolo/PROSTATEx_masks.
 38
 39This dataset is from the publications https://doi.org/10.1109/TMI.2014.2303821 (PROSTATEx)
 40https://doi.org/10.1016/j.ejrad.2021.109647 (masks) and
 41https://doi.org/10.1038/s41597-024-03945-2 (ProstateZones, CC BY 4.0).
 42The data was released at https://doi.org/10.7937/K9TCIA.2017.MURS5CL.
 43Please cite them if you use this dataset in your research.
 44"""
 45
 46import os
 47import re
 48import csv
 49import json
 50from glob import glob
 51from tqdm import tqdm
 52from natsort import natsorted
 53from collections import defaultdict
 54from typing import Union, Tuple, List, Literal
 55
 56import numpy as np
 57import requests
 58
 59from torch.utils.data import Dataset, DataLoader
 60
 61import torch_em
 62
 63from .. import util
 64
 65
 66MASKS_COMMIT = "21b9dfde9da4f7b719c206fe1ca00ae31d6f5cf3"
 67
 68URLS = {
 69    "images": f"{util.NBIA_API_URL}getSeries?Collection=PROSTATEx",
 70    "masks": f"https://github.com/rcuocolo/PROSTATEx_masks/archive/{MASKS_COMMIT}.zip",
 71    "zones_detailed": "https://zenodo.org/records/10718469/files/ProstateZones.zip?download=1",
 72    "series_list": (
 73        "https://raw.githubusercontent.com/UMU-DDI/ProstateZones/main/Support%20Files/list_of_PROSTATEx_files.txt"
 74    ),
 75}
 76
 77CHECKSUMS = {
 78    "images": None,  # The DICOM series are downloaded individually from TCIA.
 79    "masks": None,  # GitHub does not guarantee stable archive checksums.
 80    "zones_detailed": "5a45816230b4bf94e88527a00754cc0b9329ec2c3459179a3931c11cb574e80a",
 81    "series_list": None,  # GitHub does not guarantee stable raw file checksums.
 82}
 83
 84ZONE_IDS = {"peripheral_zone": 1, "transition_zone": 2}
 85"""The zone ids of the 'zones' labels."""
 86
 87DETAILED_ZONE_IDS = {
 88    "peripheral_zone": 1, "central_zone": 2, "transition_zone": 3,
 89    "anterior_fibromuscular_stroma": 4, "urethra": 5,
 90}
 91"""The zone ids of the 'zones_detailed' labels."""
 92
 93
 94def _load_dicom_volume(series_dir):
 95    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted along the slice normal."""
 96    import pydicom
 97
 98    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 99    slices = [dcm for dcm in slices if hasattr(dcm, "ImagePositionPatient")]
100    orientation = np.array([float(v) for v in slices[0].ImageOrientationPatient])
101    normal = np.cross(orientation[:3], orientation[3:])
102    slices.sort(key=lambda dcm: np.dot(np.array([float(v) for v in dcm.ImagePositionPatient]), normal))
103    return np.stack([dcm.pixel_array for dcm in slices])
104
105
106def _nifti_to_dicom_grid(data):
107    """Map a NIfTI array (x, y, z; dcm2niix conversion of an axial series) onto the DICOM grid (z, y, x).
108
109    dcm2niix flips the row axis of axial series (DICOM rows run anterior -> posterior, the NIfTI y axis runs
110    posterior -> anterior), so the array is transposed and flipped along y.
111    """
112    return np.asarray(data).transpose(2, 1, 0)[:, ::-1]
113
114
115def _load_nifti_on_dicom_grid(path):
116    import nibabel as nib
117
118    nifti = nib.load(path)
119    assert nib.aff2axcodes(nifti.affine) == ("L", "A", "S"), f"Unexpected axes for {path}"
120    return _nifti_to_dicom_grid(nifti.dataobj)
121
122
123def _parse_image_name(name):
124    """Parse a name like 'ProstateX-0000_t2_tse_tra_4' into the patient id and the DICOM series number."""
125    patient_id, _, series_number = re.match(r"(ProstateX-\d+)_(.*)_(\d+)$", name.strip()).groups()
126    return patient_id, int(series_number)
127
128
129def _read_image_lists(mask_dir):
130    """Read the lists of the DICOM series the masks were drawn on.
131
132    Returns a dict {(label_type, sequence, patient_id): (series_number, nifti_image_path)}.
133    """
134    image_series = {}
135    with open(os.path.join(mask_dir, "lesions", "Image_list.csv"), "r") as f:
136        for row in csv.DictReader(f):
137            for sequence in ("T2", "ADC"):
138                patient_id, series_number = _parse_image_name(row[sequence])
139                image_path = os.path.join(mask_dir, "lesions", "Images", sequence, f"{row[sequence].strip()}.nii.gz")
140                image_series[("lesions", sequence.lower(), patient_id)] = (series_number, image_path)
141    with open(os.path.join(mask_dir, "prostate", "image_list.csv"), "r") as f:
142        for row in csv.DictReader(f):
143            patient_id, series_number = _parse_image_name(row["T2"])
144            image_path = os.path.join(mask_dir, "prostate", "Images", f"{row['T2'].strip()}.nii.gz")
145            image_series[("zones", "t2", patient_id)] = (series_number, image_path)
146    return image_series
147
148
149def _nrrd_on_dicom_grid(path):
150    """Map a ProstateZones nrrd array (x, y, z) onto the DICOM grid (z, y, x).
151
152    These are stored with the geometry of the series they were drawn on, so unlike the dcm2niix NIfTI
153    conversions of the lesion and zone masks they only have to be transposed and not flipped.
154    """
155    import nrrd
156
157    data, _ = nrrd.read(path)
158    return np.asarray(data).transpose(2, 1, 0)
159
160
161def _read_detailed_zone_series(path):
162    """Read the axial T2 series number of every patient from the ProstateZones series list.
163
164    Each line names the series of one patient, starting with the axial T2 series, e.g.
165    'ProstateX-0000:4.000000-t2tsetra-00702:3.000000-t2tsesag-87368:...'.
166    """
167    series_numbers = {}
168    with open(os.path.join(path, "list_of_PROSTATEx_files.txt"), "r") as f:
169        for line in f:
170            fields = line.strip().split(":")
171            if len(fields) < 2:
172                continue
173            match = re.match(r"([\d.]+)-t2tsetra-", fields[1])
174            if match is not None:
175                series_numbers[fields[0]] = int(float(match.group(1)))
176    return series_numbers
177
178
179def _get_detailed_zone_masks(zones_dir, patient_id):
180    """Get the ProstateZones masks of a patient, which are one file or one file per reader."""
181    number = patient_id.split("-")[-1]
182    single = os.path.join(zones_dir, "Singles", f"Seg-{number}.nrrd")
183    if os.path.exists(single):
184        return [single]
185    duplicates = [os.path.join(zones_dir, "Duplicates", reader, f"Seg-{number}_{reader}.nrrd")
186                  for reader in ("R1", "R2")]
187    return duplicates if all(os.path.exists(mask_path) for mask_path in duplicates) else []
188
189
190def _load_detailed_zone_labels(mask_paths, shape):
191    """Load the ProstateZones masks of a patient, one per reader."""
192    labels = [_nrrd_on_dicom_grid(mask_path) for mask_path in mask_paths]
193    if any(mask.shape != shape for mask in labels):
194        return None
195    return [mask.astype("uint8") for mask in labels]
196
197
198def _get_series_metadata(path, download):
199    """Get the metadata of all series in the collection from the NBIA REST API."""
200    metadata_path = os.path.join(path, "prostatex_series.json")
201    if not os.path.exists(metadata_path):
202        if not download:
203            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
204        response = requests.get(URLS["images"])
205        response.raise_for_status()
206        with open(metadata_path, "w") as f:
207            json.dump(response.json(), f, indent=2)
208
209    with open(metadata_path, "r") as f:
210        return json.load(f)
211
212
213def _find_mask(directory, patient_id, suffix=""):
214    """Find the mask of a patient in one of the folders with the gland and zone masks.
215
216    The file names are inconsistent in the mask repository (e.g. 'ProstateX-080.nii.gz' instead of
217    'ProstateX-0080.nii.gz' and 'VOLUME-0202_pz.nii.gz' instead of 'ProstateX-0202_pz.nii.gz'),
218    so the masks are matched via the patient number rather than via the file name.
219    """
220    number = int(patient_id.split("-")[-1])
221    for mask_path in natsorted(glob(os.path.join(directory, f"*{suffix}.nii.gz"))):
222        if int(re.search(r"(\d+)", os.path.basename(mask_path)).group(1)) == number:
223            return mask_path
224    return None
225
226
227def _get_mask_paths(mask_dir, label_type, sequence, patient_id):
228    """Get the mask files of a patient for the given label type and sequence."""
229    if label_type == "lesions":
230        return natsorted(glob(os.path.join(mask_dir, "lesions", "Masks", sequence.upper(), f"{patient_id}-*")))
231
232    prostate_dir = os.path.join(mask_dir, "prostate")
233    mask_paths = [
234        _find_mask(os.path.join(prostate_dir, "mask_prostate"), patient_id),
235        _find_mask(os.path.join(prostate_dir, "mask_pz"), patient_id, "_pz"),
236        _find_mask(os.path.join(prostate_dir, "mask_tz"), patient_id, "_tz"),
237    ]
238    return mask_paths if all(mask_path is not None for mask_path in mask_paths) else []
239
240
241def _find_series(dicom_dir, candidates, image_path, mask_shape):
242    """Find the downloaded DICOM series that matches the NIfTI image the masks were drawn on.
243
244    The NIfTI images that ship with the masks are used to identify the series: the stacked DICOM volume has to
245    match the NIfTI image voxel by voxel. For the few patients without a NIfTI image in the mask repository the
246    series is identified by the shape of the masks instead.
247
248    Returns the volume on the DICOM grid or None if no candidate series matches.
249    """
250    reference = _load_nifti_on_dicom_grid(image_path) if os.path.exists(image_path) else None
251    for series in candidates:
252        series_dir = os.path.join(dicom_dir, series["SeriesInstanceUID"])
253        if not os.path.exists(series_dir):
254            continue
255        volume = _load_dicom_volume(series_dir)
256        if reference is None:
257            if volume.shape == mask_shape:
258                return volume
259        elif volume.shape == reference.shape and np.array_equal(volume, reference.astype(volume.dtype)):
260            return volume
261    return None
262
263
264def _load_lesion_labels(mask_paths, patient_id, shape):
265    """Combine the per-lesion masks of a patient into an instance segmentation (id = finding id)."""
266    labels = np.zeros(shape, dtype="uint8")
267    pattern = re.compile(rf"{patient_id}-Finding(\d+)-.*ROI\.nii\.gz", re.IGNORECASE)
268    finding_ids = [int(pattern.match(os.path.basename(mask_path)).group(1)) for mask_path in mask_paths]
269    n_masks = 0
270    for mask_path, finding_id in zip(mask_paths, finding_ids):
271        # The finding ids start at 1, except for a single lesion in the current release of the masks
272        # (ProstateX-0005), which is called 'Finding0' and gets the next free id instead of the background id.
273        if finding_id == 0:
274            finding_id = max(finding_ids) + 1
275        mask = _load_nifti_on_dicom_grid(mask_path)
276        if mask.shape != shape:  # A few masks were drawn on a different series and cannot be used.
277            continue
278        labels[mask > 0] = finding_id
279        n_masks += 1
280    return labels if n_masks > 0 else None
281
282
283def _load_zone_labels(mask_paths, shape):
284    """Combine the peripheral and transition zone masks into a semantic segmentation."""
285    prostate, pz, tz = [_load_nifti_on_dicom_grid(mask_path) > 0 for mask_path in mask_paths]
286    if any(mask.shape != shape for mask in (prostate, pz, tz)):
287        return None, None
288    zones = np.zeros(shape, dtype="uint8")
289    zones[tz] = ZONE_IDS["transition_zone"]
290    zones[pz] = ZONE_IDS["peripheral_zone"]
291    return zones, prostate.astype("uint8")
292
293
294def _preprocess_prostatex(dicom_dir, mask_dir, zones_dir, series_metadata, image_series, preprocessed_dir):
295    import h5py
296
297    series_by_number = defaultdict(list)
298    for series in series_metadata:
299        series_by_number[(series["PatientID"], int(series["SeriesNumber"]))].append(series)
300
301    groups_per_patient = defaultdict(list)
302    for (label_type, sequence, patient_id), (series_number, image_path) in image_series.items():
303        groups_per_patient[patient_id].append((label_type, sequence, series_number, image_path))
304
305    os.makedirs(preprocessed_dir, exist_ok=True)
306    for patient_id, groups in tqdm(sorted(groups_per_patient.items()), desc="Preprocess PROSTATEx"):
307        out_path = os.path.join(preprocessed_dir, f"{patient_id}.h5")
308        # Groups that are already stored are kept, so that a file written by an earlier version of this
309        # module gains the groups it is missing instead of being recomputed.
310        if os.path.exists(out_path):
311            with h5py.File(out_path, "r") as f:
312                stored = {key for key in f if isinstance(f[key], h5py.Group)}
313            groups = [group for group in groups if group[0] not in stored]
314            if not groups:
315                continue
316
317        datasets = {}
318        for label_type, sequence, series_number, image_path in groups:
319            if label_type == "zones_detailed":
320                mask_paths = _get_detailed_zone_masks(zones_dir, patient_id)
321            else:
322                mask_paths = _get_mask_paths(mask_dir, label_type, sequence, patient_id)
323            if not mask_paths:  # Some patients only have masks for one of the label types.
324                continue
325            if label_type == "zones_detailed":
326                mask_shape = _nrrd_on_dicom_grid(mask_paths[0]).shape
327            else:
328                mask_shape = _load_nifti_on_dicom_grid(mask_paths[0]).shape
329            candidates = series_by_number[(patient_id, series_number)]
330            volume = _find_series(dicom_dir, candidates, image_path, mask_shape)
331            if volume is None:
332                raise RuntimeError(f"No DICOM series matches the masks of {patient_id} ({label_type}, {sequence}).")
333            if label_type == "lesions":
334                labels = _load_lesion_labels(mask_paths, patient_id, volume.shape)
335                if labels is None:
336                    continue
337            elif label_type == "zones_detailed":
338                reader_labels = _load_detailed_zone_labels(mask_paths, volume.shape)
339                if reader_labels is None:
340                    continue
341                labels = reader_labels[0]
342                # 40 of the patients were delineated by two readers independently.
343                if len(reader_labels) > 1:
344                    datasets[f"{label_type}/{sequence}/labels_reader2"] = reader_labels[1]
345            else:
346                labels, prostate = _load_zone_labels(mask_paths, volume.shape)
347                if labels is None:
348                    continue
349                datasets[f"{label_type}/{sequence}/prostate"] = prostate
350            datasets[f"{label_type}/{sequence}/raw"] = volume
351            datasets[f"{label_type}/{sequence}/labels"] = labels
352
353        if not datasets:
354            continue
355        if os.path.exists(out_path):
356            with h5py.File(out_path, "a") as f:
357                for key, data in datasets.items():
358                    f.create_dataset(key, data=data, compression="gzip")
359        else:
360            tmp_path = out_path + ".tmp"
361            with h5py.File(tmp_path, "w") as f:
362                for key, data in datasets.items():
363                    f.create_dataset(key, data=data, compression="gzip")
364            os.rename(tmp_path, out_path)
365
366
367def get_prostatex_data(path: Union[os.PathLike, str], download: bool = False) -> str:
368    """Download the PROSTATEx dataset.
369
370    The download is resumable: series that were already downloaded and patients that were already converted are
371    skipped when the function is called again.
372
373    Args:
374        path: Filepath to a folder where the data is downloaded for further processing.
375        download: Whether to download the data if it is not present.
376
377    Returns:
378        Filepath where the preprocessed data is stored.
379    """
380    os.makedirs(path, exist_ok=True)
381    preprocessed_dir = os.path.join(path, "preprocessed")
382
383    # Download the masks.
384    mask_dir = os.path.join(path, f"PROSTATEx_masks-{MASKS_COMMIT}", "Files")
385    if not os.path.exists(mask_dir):
386        zip_path = os.path.join(path, f"PROSTATEx_masks-{MASKS_COMMIT}.zip")
387        util.download_source(path=zip_path, url=URLS["masks"], download=download, checksum=CHECKSUMS["masks"])
388        util.unzip(zip_path=zip_path, dst=path)
389
390    # Download the ProstateZones annotations and the list of the series they were drawn on.
391    zones_dir = os.path.join(path, "ProstateZones")
392    if not os.path.exists(os.path.join(zones_dir, "Singles")):
393        zip_path = os.path.join(path, "ProstateZones.zip")
394        util.download_source(
395            path=zip_path, url=URLS["zones_detailed"], download=download, checksum=CHECKSUMS["zones_detailed"]
396        )
397        util.unzip(zip_path=zip_path, dst=zones_dir, remove=False)
398    util.download_source(
399        path=os.path.join(path, "list_of_PROSTATEx_files.txt"), url=URLS["series_list"],
400        download=download, checksum=CHECKSUMS["series_list"],
401    )
402
403    image_series = _read_image_lists(mask_dir)
404    detailed_series = _read_detailed_zone_series(path)
405    for patient_id, series_number in detailed_series.items():
406        if _get_detailed_zone_masks(zones_dir, patient_id):
407            # The ProstateZones masks ship without a reference image, so the series is matched by shape.
408            image_series[("zones_detailed", "t2", patient_id)] = (series_number, "")
409
410    n_patients = len({patient_id for _, _, patient_id in image_series})
411    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == n_patients:
412        return preprocessed_dir
413
414    # Find the series the masks were drawn on and download them from TCIA.
415    series_metadata = _get_series_metadata(path, download)
416    series_numbers = {(patient_id, number) for (_, _, patient_id), (number, _) in image_series.items()}
417    series_uids = sorted(
418        series["SeriesInstanceUID"] for series in series_metadata
419        if (series["PatientID"], int(series["SeriesNumber"])) in series_numbers
420    )
421    dicom_dir = os.path.join(path, "dicom")
422    if download:  # Series that were downloaded already are skipped.
423        util.download_tcia_series(series_uids, dst=dicom_dir, csv_filename=os.path.join(path, "prostatex_series"))
424    elif not all(os.path.exists(os.path.join(dicom_dir, uid)) for uid in series_uids):
425        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
426
427    _preprocess_prostatex(dicom_dir, mask_dir, zones_dir, series_metadata, image_series, preprocessed_dir)
428    return preprocessed_dir
429
430
431def get_prostatex_paths(
432    path: Union[os.PathLike, str],
433    sequence: Literal["t2", "adc"] = "t2",
434    label_type: Literal["lesions", "zones", "zones_detailed"] = "lesions",
435    download: bool = False,
436) -> List[str]:
437    """Get paths to the PROSTATEx data.
438
439    Args:
440        path: Filepath to a folder where the data is downloaded for further processing.
441        sequence: The MRI sequence, either 't2' or 'adc'. The zone labels are only available for 't2'.
442        label_type: The label type, one of 'lesions', 'zones' or 'zones_detailed'.
443        download: Whether to download the data if it is not present.
444
445    Returns:
446        List of filepaths for the hdf5 files, which contain the image data ('<label_type>/<sequence>/raw')
447        and the label data ('<label_type>/<sequence>/labels').
448    """
449    import h5py
450
451    assert sequence in ("t2", "adc"), f"Invalid sequence: {sequence}."
452    assert label_type in ("lesions", "zones", "zones_detailed"), f"Invalid label type: {label_type}."
453    if label_type in ("zones", "zones_detailed") and sequence != "t2":
454        raise ValueError("The zone labels are only available for the 't2' sequence.")
455
456    data_dir = get_prostatex_data(path, download)
457    volume_paths = []
458    for volume_path in natsorted(glob(os.path.join(data_dir, "*.h5"))):
459        with h5py.File(volume_path, "r") as f:
460            if f"{label_type}/{sequence}/labels" in f:
461                volume_paths.append(volume_path)
462    return volume_paths
463
464
465def get_prostatex_dataset(
466    path: Union[os.PathLike, str],
467    patch_shape: Tuple[int, ...],
468    sequence: Literal["t2", "adc"] = "t2",
469    label_type: Literal["lesions", "zones", "zones_detailed"] = "lesions",
470    resize_inputs: bool = False,
471    download: bool = False,
472    **kwargs
473) -> Dataset:
474    """Get the PROSTATEx dataset for prostate lesion or zone segmentation.
475
476    Args:
477        path: Filepath to a folder where the data is downloaded for further processing.
478        patch_shape: The patch shape to use for training.
479        sequence: The MRI sequence, either 't2' or 'adc'. The zone labels are only available for 't2'.
480        label_type: The label type, one of 'lesions', 'zones' or 'zones_detailed'.
481        resize_inputs: Whether to resize inputs to the desired patch shape.
482        download: Whether to download the data if it is not present.
483        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
484
485    Returns:
486        The segmentation dataset.
487    """
488    volume_paths = get_prostatex_paths(path, sequence, label_type, download)
489
490    if resize_inputs:
491        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
492        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
493            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
494        )
495
496    return torch_em.default_segmentation_dataset(
497        raw_paths=volume_paths,
498        raw_key=f"{label_type}/{sequence}/raw",
499        label_paths=volume_paths,
500        label_key=f"{label_type}/{sequence}/labels",
501        patch_shape=patch_shape,
502        is_seg_dataset=True,
503        **kwargs
504    )
505
506
507def get_prostatex_loader(
508    path: Union[os.PathLike, str],
509    batch_size: int,
510    patch_shape: Tuple[int, ...],
511    sequence: Literal["t2", "adc"] = "t2",
512    label_type: Literal["lesions", "zones", "zones_detailed"] = "lesions",
513    resize_inputs: bool = False,
514    download: bool = False,
515    **kwargs
516) -> DataLoader:
517    """Get the PROSTATEx dataloader for prostate lesion or zone segmentation.
518
519    Args:
520        path: Filepath to a folder where the data is downloaded for further processing.
521        batch_size: The batch size for training.
522        patch_shape: The patch shape to use for training.
523        sequence: The MRI sequence, either 't2' or 'adc'. The zone labels are only available for 't2'.
524        label_type: The label type, one of 'lesions', 'zones' or 'zones_detailed'.
525        resize_inputs: Whether to resize inputs to the desired patch shape.
526        download: Whether to download the data if it is not present.
527        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
528
529    Returns:
530        The DataLoader.
531    """
532    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
533    dataset = get_prostatex_dataset(path, patch_shape, sequence, label_type, resize_inputs, download, **ds_kwargs)
534    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)