DicomConvert/DicomConvert.py

# =========================================================================
#
#  Copyright NumFOCUS
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0.txt
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
# =========================================================================
 
import argparse
import csv
import functools
import itertools
import multiprocessing
import os
import sys
 
import SimpleITK as sitk
 
 
def convert_image(input_file_name, output_file_name, new_width=None):
    try:
        image_file_reader = sitk.ImageFileReader()
        # only read DICOM images
        image_file_reader.SetImageIO("GDCMImageIO")
        image_file_reader.SetFileName(input_file_name)
        image_file_reader.ReadImageInformation()
        image_size = list(image_file_reader.GetSize())
        if len(image_size) == 3 and image_size[2] == 1:
            image_size[2] = 0
        image_file_reader.SetExtractSize(image_size)
        image = image_file_reader.Execute()
        if new_width:
            original_size = image.GetSize()
            original_spacing = image.GetSpacing()
            new_spacing = [
                (original_size[0] - 1) * original_spacing[0] / (new_width - 1)
            ] * 2
            new_size = [
                new_width,
                int(
                    (original_size[1] - 1)
                    * original_spacing[1]
                    / new_spacing[1]
                ),
            ]
            image = sitk.Resample(
                image1=image,
                size=new_size,
                transform=sitk.Transform(),
                interpolator=sitk.sitkLinear,
                outputOrigin=image.GetOrigin(),
                outputSpacing=new_spacing,
                outputDirection=image.GetDirection(),
                defaultPixelValue=0,
                outputPixelType=image.GetPixelID(),
            )
        # If a single channel image, rescale to [0,255]. Also modify the
        # intensity values based on the photometric interpretation. If
        # MONOCHROME2 (minimum should be displayed as black) we don't need to
        # do anything, if image has MONOCRHOME1 (minimum should be displayed as
        # white) we flip # the intensities. This is a constraint imposed by ITK
        # which always assumes MONOCHROME2.
        if image.GetNumberOfComponentsPerPixel() == 1:
            image = sitk.RescaleIntensity(image, 0, 255)
            if (
                image_file_reader.GetMetaData("0028|0004").strip()
                == "MONOCHROME1"
            ):
                image = sitk.InvertIntensity(image, maximum=255)
            image = sitk.Cast(image, sitk.sitkUInt8)
        sitk.WriteImage(image, output_file_name)
        return True
    except BaseException:
        return False
 
 
def convert_images(input_file_names, output_file_names, new_width):
    MAX_PROCESSES = 15
    with multiprocessing.Pool(processes=MAX_PROCESSES) as pool:
        return pool.starmap(
            functools.partial(convert_image, new_width=new_width),
            zip(input_file_names, output_file_names),
        )
 
 
def positive_int(int_str):
    value = int(int_str)
    if value <= 0:
        raise argparse.ArgumentTypeError(
            int_str + " is not a positive integer value"
        )
    return value
 
 
def directory(dir_name):
    if not os.path.isdir(dir_name):
        raise argparse.ArgumentTypeError(
            dir_name + " is not a valid directory name"
        )
    return dir_name
 
 
def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Convert and resize DICOM files to common image types."
    )
    parser.add_argument(
        "root_of_data_directory",
        type=directory,
        help="Path to the topmost directory containing data.",
    )
    parser.add_argument(
        "output_file_extension",
        help="Image file extension, this determines output file type "
        "(e.g. png) .",
    )
    parser.add_argument(
        "--w", type=positive_int, help="Width of converted images."
    )
    parser.add_argument("--od", type=directory, help="Output directory.")
    args = parser.parse_args(argv)
 
    input_file_names = []
    for dir_name, subdir_names, file_names in os.walk(
        args.root_of_data_directory
    ):
        input_file_names += [
            os.path.join(os.path.abspath(dir_name), fname)
            for fname in file_names
        ]
    if args.od:
        # if all output files are written to the same directory we need them
        # to have a unique name, so use an index.
        file_names = [
            os.path.join(os.path.abspath(args.od), str(i))
            for i in range(len(input_file_names))
        ]
    else:
        file_names = input_file_names
    output_file_names = [
        file_name + "." + args.output_file_extension
        for file_name in file_names
    ]
 
    res = convert_images(input_file_names, output_file_names, args.w)
    input_file_names = list(itertools.compress(input_file_names, res))
    output_file_names = list(itertools.compress(output_file_names, res))
 
    # save csv file mapping input and output file names.
    # using csv module and not pandas so as not to create more dependencies
    # for the examples. pandas based code is more elegant/shorter.
    dir_name = args.od if args.od else os.getcwd()
    with open(os.path.join(dir_name, "file_names.csv"), mode="w") as fp:
        fp_writer = csv.writer(
            fp, delimiter=",", quotechar='"', quoting=csv.QUOTE_MINIMAL
        )
        fp_writer.writerow(["input file name", "output file name"])
        for data in zip(input_file_names, output_file_names):
            fp_writer.writerow(data)
 
 
if __name__ == "__main__":
    sys.exit(main())