dicom-utilities/dataladify_dicom_dataset.py

#!/usr/bin/env python3
from datetime import datetime
from operator import itemgetter
from pathlib import Path, PurePath
import tarfile

import click
from datalad.api import create, ls_file_collection
from pydicom import dcmread
from pydicom.valuerep import DA, TM
from pydicom.errors import InvalidDicomError
from rich.progress import track

default_date = datetime(1970, 1, 1)
dicom_metadata_keys = [
    "SeriesDescription",
    "SeriesNumber",
    "Modality",
    "MRAcquisitionType",
    "ProtocolName",
    "PulseSequenceName",
]


def dicom_datetime(p: Path):
    """Return a datetime object combining StudyDate & StudyTime

    Returns None if the file under path is not a DICOM
    """
    try:
        with dcmread(p, stop_before_pixels=True) as dcm:
            studydate = DA(dcm.StudyDate)
            studytime = TM(dcm.StudyTime)
            study_dt = datetime.combine(studydate, studytime)
    except InvalidDicomError:
        study_dt = None
    return study_dt


def normalize_tarinfo(tinfo, mtime_dt):
    """Standardize permissions, ownership, and timestamp"""
    # be safe
    tinfo.uid = 0
    tinfo.gid = 0
    tinfo.uname = "root"
    tinfo.gname = "root"
    tinfo.mtime = mtime_dt.timestamp()
    if tinfo.isfile():
        # for any regular file normalize the permission
        # leave unexpected extra-ordinary content untouched
        tinfo.mode = int("0o100644", 8)
    return tinfo


def generate_tarball(tar_path, fc):
    """Create a tarball with all annexed files"""
    with tarfile.open(tar_path, "w", dereference=True) as tar:
        for element in track(fc, description="Building tar file"):
            if element["type"] != "annexed file":
                continue

            # item is a PurePosixPath relative to the collection
            # without checksumming ls_collection gives no fp
            # so we create a Path object to allow access
            element_path = element["collection"] / element["item"]

            dt = dicom_datetime(element_path)
            if dt is None:
                dt = default_date

            tinfo = tar.gettarinfo(name=element_path, arcname=element["item"])
            tinfo = normalize_tarinfo(tinfo, dt)

            with element_path.open("rb") as fp:
                tar.addfile(tinfo, fp)


def extract_dicom_metadata(fc, dicom_metadata_keys):
    """Extract information for addurls, including selected dicom keys"""
    dicoms = []
    for element in track(fc, description="Collecting DICOM file information"):
        if element["type"] != "annexed file":
            continue

        # path object that can be opened
        element_path = element["collection"] / element["item"]

        d = {
            "path": element["item"],
            "size": element["annexsize"],
            "annexkey": element["annexkey"],
        }

        try:
            with dcmread(element_path, stop_before_pixels=True) as dcm:
                d.update({dmk: getattr(dcm, dmk, "") for dmk in dicom_metadata_keys})
        except InvalidDicomError:
            pass

        dicoms.append(d)

    return dicoms


@click.command()
@click.argument(
    "dicomdir",
    type=click.Path(exists=True, file_okay=False, dir_okay=True, path_type=Path),
)
def main(dicomdir):
    """Dataladify DICOMDIR in-place

    This utility will create a DataLad dataset in the selected folder,
    pack all contents into a tar file, and record file availability in
    the tar archive. Contents other than the archive itself will then
    be dropped.

    Recording the file availability in this way means that only one
    annex key needs to be stored, therefore saving inodes on file
    systems where this is important. The Git repository can be packed
    with the standard Git means.

    The tar archive is generated with standardized file order,
    modification time, permissions, and ownership to yield the same
    checksum reproducibly. Files in the tar archive are declared to be
    owned by root with 0644 permissions. Modification date for each
    DICOM file is determined by its DICOM StudyDate and StudyTime
    timestamps.

    """

    # name the tarball using the directory name, so it can be used standalone
    # put it under "archives"
    tar_relpath = PurePath("archives", f"{dicomdir.name}.tar")
    tar_path = dicomdir / tar_relpath

    # be safe and avoid messing with existing datasets
    assert not tar_path.exists()
    assert not (dicomdir / ".datalad").exists()
    tar_path.parent.mkdir(exist_ok=True)

    # create the dataset in-place and save all files
    ds = create(dicomdir, force=True)
    ds.save(message="Save dicoms")

    # List file collection & sort in-place for deterministic tarball order
    fc = ls_file_collection("annexworktree", dicomdir, result_renderer="disabled")
    fc.sort(key=itemgetter("item"))

    # save the tarball and learn about its annex key
    generate_tarball(tar_path, fc)
    res = ds.save(path=tar_relpath, message="Save DICOM tarball")
    added_keys = [r.get("key") for r in res if r.get("action") == "add"]
    assert len(added_keys) == 1
    tarkey = added_keys[0]

    # register the archivist special remote, to claim the dl+archives URLs registered below.
    ds.repo.call_annex(
        [
            "initremote",
            "archivist",
            "type=external",
            "externaltype=archivist",
            "encryption=none",
            # auto-enabling is cheap (makes no connection attempts), and convenient
            "autoenable=true",
        ]
    )
    archivist_uuid = ds.repo.call_annex_records(["info", "archivist"])[0]["uuid"]
    assert archivist_uuid

    # add dl+archives URLs to all files; also add selected DICOM metadata
    dicoms = extract_dicom_metadata(fc, dicom_metadata_keys)
    dicom_recs = ds.addurls(
        dicoms,
        urlformat=f"dl+archive:{tarkey}#path={{path}}&size={{size}}",
        filenameformat="{path}",
        key="{annexkey}",
        exclude_autometa="path|size|annexkey",
    )

    # assure availability for each DICOM
    dicomkeys = [r["annexkey"] for r in dicom_recs if r.get("action") == "fromkey"]
    for dicomkey in dicomkeys:
        ds.repo.call_annex(["setpresentkey", dicomkey, archivist_uuid, "1"])

    # probe the availability metadata. This seems to be necessary at times to
    # get git-annex to commit the metadata operations performed above
    # to be able to actually push everything
    ds.repo.call_annex(["whereis", "--key", dicomkeys[0]])

    # drop the dicom files, leaving just the tar file behind
    ds.repo.call_annex(["drop", ".", "--exclude", "*.tar"])


if __name__ == "__main__":
    main()