add new preparation script

mu_map/data/prepare2.py

+"""
+New dataset preparation script because the way DICOM data is exported has changed.
+"""
+import argparse
+from enum import Enum
+import os
+from typing import Any, Dict, List
+
+import pandas as pd
+import pydicom
+
+from mu_map.data.prepare import headers
+from mu_map.file.dicom import DICOMTime, EXTENSION_DICOM, parse_age
+from mu_map.logging import add_logging_args, get_logger_by_args
+
+
+class Protocol(Enum):
+    """
+    Class defining the two protocols __Rest__ & __Stress__ of MPI
+    SPECT studies.
+    """
+
+    Rest = 1
+    Stress = 2
+
+    @classmethod
+    def get(cls, dcm: pydicom.dataset.Dataset):
+        """
+        Get the protocol from a DICOM dataset.
+
+        This function is based on the study description.
+
+        Parameters
+        ----------
+        dcm: pydicom.dataset.Dataset
+            the DICOM dataset
+
+        Returns
+        -------
+        Protocol
+        """
+        if Protocol.Rest.name.lower() in dcm.SeriesDescription.lower():
+            return Protocol.Rest
+        elif Protocol.Stress.name.lower() in dcm.SeriesDescription.lower():
+            return Protocol.Stress
+        else:
+            raise ValueError(
+                f"Cannot extract protocol from SeriesDescription[{dcm.SeriesDescription}]"
+            )
+
+
+def is_scatter_corrected(dcm: pydicom.dataset.Dataset) -> bool:
+    """
+    Test if a DICOM dataset is a scatter-corrected reconstrucion based
+    on the __CorrectedImage__ tag.
+
+    Parameters
+    ----------
+    dcm: pydicom.dataset.Dataset
+        the DICOM dataset
+
+    Returns
+    -------
+    bool
+    """
+    try:
+        return "SCAT" in dcm.CorrectedImage
+    except AttributeError:
+        return False
+
+
+def is_attenuation_corrected(dcm: pydicom.dataset.Dataset) -> bool:
+    """
+    Test if a DICOM dataset is an attenuation-corrected reconstrucion based
+    on the __CorrectedImage__ tag.
+
+    Parameters
+    ----------
+    dcm: pydicom.dataset.Dataset
+        the DICOM dataset
+
+    Returns
+    -------
+    bool
+    """
+    try:
+        return "ATTN" in dcm.CorrectedImage
+    except AttributeError:
+        return False
+
+
+def is_mu_map(dcm: pydicom.dataset.Dataset) -> bool:
+    """
+    Test if a DICOM dataset is a mu map based on the
+    series description.
+
+    Parameters
+    ----------
+    dcm: pydicom.dataset.Dataset
+        the DICOM dataset
+
+    Returns
+    -------
+    bool
+    """
+    return "µ-map" in dcm.SeriesDescription
+
+
+def closest_in_time(
+    ref: pydicom.dataset.Dataset,
+    dcms: List[pydicom.dataset.Dataset],
+    time_type: DICOMTime = DICOMTime.Series,
+) -> pydicom.dataset.Dataset:
+    """
+    Find the DICOM file closest in time to a reference file.
+
+    Parameters
+    ----------
+    ref: pydicom.dataset.Dataset
+        the reference DICOM file
+    dcms: List[pydicom.dataset.Dataset]
+        the list which is searched
+    time_type: DICOMTime
+        the time type used for the search
+
+    Returns
+    -------
+    pydicom.dataset.Dataset
+    """
+    ref_time = time_type.to_datetime(ref)
+    dcms.sort(key=lambda dcm: abs(time_type.to_datetime(dcm) - ref_time))
+    return dcms[0]
+
+
+def get_recon_filename(dcm: pydicom.dataset.Dataset) -> str:
+    """
+    Create a filename for a reconstruction depending on if it is
+    scatter- and/or attenuation-corrected.
+
+    Parameters
+    ----------
+    dcm: pydicom.dataset.Dataset
+        the dataset for which a filename is created
+
+    Returns
+    -------
+    str
+    """
+    sc_str = "sc" if is_scatter_corrected(dcm) else "nsc"
+    ac_str = "ac" if is_attenuation_corrected(dcm) else "nac"
+    return f"recon_{ac_str}_{sc_str}"
+
+
+def create_row(
+    _id: int, dcm: pydicom.dataset.Dataset, protocol: Protocol
+) -> Dict[str, Any]:
+    """
+    Create a row for the meta data table.
+
+    Parameters
+    ----------
+    _id: int
+        the id of the row
+    dcm: pydicom.dataset.Dataset
+        the DICOM file to extract meta information
+    protocol: Protocol
+        the protocol of the study
+
+    Returns
+    -------
+    Dict[str, Any]
+    """
+    return {
+        headers.id: _id,
+        headers.patient_id: dcm.PatientID,
+        headers.age: parse_age(dcm.PatientAge),
+        headers.sex: dcm.PatientSex,
+        headers.weight: float(dcm.PatientWeight),
+        headers.size: float(dcm.PatientSize),
+        headers.protocol: protocol.name.lower(),
+        headers.datetime_acquisition: DICOMTime.Acquisition.to_datetime(dcm),
+    }
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Prepare a dataset from DICOM directories",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--dicom_dir",
+        type=str,
+        required=True,
+        help="directory containing the raw DICOM files to be processed - it is expected to contain directories and each directory can contain multiple studies of the same patient",
+    )
+    parser.add_argument(
+        "--dataset_dir",
+        type=str,
+        required=True,
+        help="directory where images, meta-information and the logs are stored",
+    )
+    parser.add_argument(
+        "--images_dir",
+        type=str,
+        default="images",
+        help="sub-directory of --dataset_dir where images are stored",
+    )
+    parser.add_argument(
+        "--meta_csv",
+        type=str,
+        default="meta.csv",
+        help="CSV file under --dataset_dir where meta-information is stored",
+    )
+    add_logging_args(
+        parser, defaults={"--logfile": "prepare.log", "--loglevel": "DEBUG"}
+    )
+    args = parser.parse_args()
+    args.images_dir = os.path.join(args.dataset_dir, args.images_dir)
+    args.meta_csv = os.path.join(args.dataset_dir, args.meta_csv)
+    args.logfile = os.path.join(args.dataset_dir, args.logfile)
+
+    if not os.path.exists(args.dataset_dir):
+        os.mkdir(args.dataset_dir)
+
+    if not os.path.exists(args.images_dir):
+        os.mkdir(args.images_dir)
+
+    global logger
+    logger = get_logger_by_args(args)
+
+    patient_dirs = sorted(os.listdir(args.dicom_dir))
+    data = pd.DataFrame({})
+    current_id = 1
+    for patient_dir in patient_dirs:
+        logger.info(f"Process directory: {patient_dir}")
+        patient_dir = os.path.join(args.dicom_dir, patient_dir)
+
+        files = sorted([os.path.join(patient_dir, f) for f in os.listdir(patient_dir)])
+        dcms = list(map(lambda f: pydicom.dcmread(f), files))
+        files = dict(zip(map(lambda dcm: dcm.SeriesInstanceUID, dcms), files))
+
+        for protocol in Protocol:
+            _dcms = list(filter(lambda dcm: Protocol.get(dcm) is protocol, dcms))
+            dcms_sc_ac = list(
+                filter(
+                    lambda dcm: is_scatter_corrected(dcm)
+                    and is_attenuation_corrected(dcm),
+                    _dcms,
+                )
+            )
+            if len(dcms_sc_ac) == 0:
+                logger.info(f" - Protocol {protocol.name} not available!")
+                continue
+
+            logger.info(f" - Process protocol: {protocol.name}")
+
+            for dcm_sc_ac in dcms_sc_ac:
+                logger.info(f"  - Create row with id: {current_id:04d}")
+                dcm_sc_nac = closest_in_time(
+                    dcm_sc_ac,
+                    list(
+                        filter(
+                            lambda dcm: is_scatter_corrected(dcm)
+                            and not is_attenuation_corrected(dcm),
+                            _dcms,
+                        )
+                    ),
+                )
+                dcm_nsc_ac = closest_in_time(
+                    dcm_sc_ac,
+                    list(
+                        filter(
+                            lambda dcm: not is_scatter_corrected(dcm)
+                            and is_attenuation_corrected(dcm),
+                            _dcms,
+                        )
+                    ),
+                )
+                dcm_nsc_nac = closest_in_time(
+                    dcm_sc_ac,
+                    list(
+                        filter(
+                            lambda dcm: not is_scatter_corrected(dcm)
+                            and not is_attenuation_corrected(dcm),
+                            _dcms,
+                        )
+                    ),
+                )
+                mu_map = closest_in_time(
+                    dcm_sc_ac, list(filter(lambda dcm: is_mu_map(dcm), _dcms))
+                )
+
+                row = create_row(current_id, dcm_sc_ac, protocol)
+
+                recons = [dcm_sc_ac, dcm_sc_nac, dcm_nsc_ac, dcm_nsc_nac]
+                _files = list(map(lambda recon: get_recon_filename(recon), recons))
+                for recon, f in zip(recons, _files):
+                    _f = f"{current_id:04d}-{f}{EXTENSION_DICOM}"
+                    row[f"file_{f}"] = _f
+                    _f = os.path.join(args.images_dir, _f)
+                    logger.info(f"  - Store {files[recon.SeriesInstanceUID]} at {_f}")
+                    pydicom.dcmwrite(_f, recon)
+
+                _f = f"{current_id:04d}-mu_map{EXTENSION_DICOM}"
+                row[headers.file_mu_map] = _f
+                _f = os.path.join(args.images_dir, _f)
+                logger.info(f"  - Store {files[mu_map.SeriesInstanceUID]} at {_f}")
+                pydicom.dcmwrite(_f, mu_map)
+
+                row = pd.DataFrame(row, index=[0])
+                data = pd.concat((data, row), ignore_index=True)
+
+                current_id = current_id + 1
+    data.to_csv(args.meta_csv, index=False)