dicoms_to_volume.py

import json
import os
from pathlib import Path
from shutil import copyfile
from tempfile import TemporaryDirectory
from typing import NamedTuple

import SimpleITK as sitk


def file_list_from_directory(directory, extension=".dcm"):
    """Return a list of files with the given extension in the given directory."""
    return [str(f) for f in Path(directory).glob(f"*{extension}")]


def write_image_and_metadata(image, metadata, output_image_filename):
    """Write the image to a file, along with the metadata."""
    sitk.WriteImage(image, output_image_filename)
    p = Path(output_image_filename)
    metadata_filename = p.parent / (p.stem + "_SOPInstanceUIDs.json")
    with open(metadata_filename, "w") as f:
        json.dump(metadata, f, indent=2)


class ImageAndMetadata(NamedTuple):
    image: sitk.Image
    metadata: dict


def metadata_dict_to_sop_instance_uids(metadata_dict):
    """Convert a metadata dictionary to a dictionary mapping slice index to SOPInstanceUID."""
    return {int(k): v["SOPInstanceUID"] for k, v in metadata_dict.items()}


def dicoms_to_volume(valid_dcm_file_list) -> ImageAndMetadata:
    """Convert a list of DICOM files to a image volume. Also returns metadata
    (SOPInstanceUID) for each slice in the volume.

    The metadata generated by this function is a map from slice index to
    SOPInstanceUID. And can be used to upload annotations from non-dicom formats
    back to MDai, which requires the SOPInstanceUID to identify the slice.

    PRECONDTIION: the input dicom files are in the same series, they are
    valid DICOM files, with no SCOUTS or other special slices.
    """
    # copy to temp directory in case multiple series occupy the same directory
    with TemporaryDirectory() as temp_dir:
        src_file_lookup = {}
        for src_file in valid_dcm_file_list:
            dst_file = os.path.join(temp_dir, os.path.basename(src_file))
            copyfile(src_file, dst_file)
            src_file_lookup[dst_file] = src_file
        reader = sitk.ImageSeriesReader()
        # allow for reading of metadata
        reader.SetMetaDataDictionaryArrayUpdate(True)
        dicom_names = reader.GetGDCMSeriesFileNames(temp_dir)
        reader.SetFileNames(dicom_names)
        image = reader.Execute()
        uids = [reader.GetMetaData(idx, "0008|0018") for idx in range(len(dicom_names))]
        # keep info about slice, original file, and SOPInstanceUID
        metadata = {
            slice_idx: {"dicom_file": src_file_lookup[fn], "SOPInstanceUID": uid}
            for slice_idx, (fn, uid) in enumerate(zip(dicom_names, uids))
        }
        return ImageAndMetadata(image=image, metadata=metadata)


def _get_parser():
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-i",
        "--input_directory",
        type=str,
        required=True,
        help="Path to the directory containing the DICOM files.",
    )
    parser.add_argument(
        "-o",
        "--output_image",
        type=str,
        required=True,
        help="Path to the output image.",
    )
    return parser


def main(input_directory, output_image):
    """
    Convert a directory containing DICOM files to a volume image, and save the metadata to map SOPInstanceUID to slice indices.
    """
    valid_dcm_file_list = file_list_from_directory(input_directory)
    image_and_metadata = dicoms_to_volume(valid_dcm_file_list)
    write_image_and_metadata(
        image=image_and_metadata.image,
        metadata=image_and_metadata.metadata,
        output_image_filename=output_image,
    )


if __name__ == "__main__":
    parser = _get_parser()
    args = parser.parse_args()
    main(args.input_directory, args.output_image)