Generate resampled MIP image from dicom

generate_mip.py

import SimpleITK as sitk
import numpy as np
from imageio import imwrite
from argparse import ArgumentParser

"""
requirements.txt

SimpleITK
numpy
scipy
imageio

example usage:
    python extract_mip.py --dicom_dir /path/to/dicom_dir
"""


def normalise_zero_one(image, eps=1e-8):
    image = image.astype(np.float32)
    ret = image - np.min(image)
    ret /= np.max(image) - np.min(image) + eps

    return ret


def reduce_hu_intensity_range(img, minv=100, maxv=1500):
    img = np.clip(img, minv, maxv)
    img = 255 * normalise_zero_one(img)

    return img


def preprocess_mip_for_slice_detection(
    image,
    spacing,
    target_spacing,
):
    # image = zoom(image, [spacing[2] / target_spacing, spacing[0] / target_spacing])
    image = reduce_hu_intensity_range(image)

    return image


def preprocess_sitk_image_for_slice_detection(
    sitk_image,
    target_spacing=1.0,
    mode="frontal",
):
    spacing = sitk_image.GetSpacing()
    direction = sitk_image.GetDirection()
    dx = int(direction[0])
    dy = int(direction[4])
    dz = int(direction[8])

    image = sitk.GetArrayFromImage(sitk_image)[::dx, ::dy, ::dz]

    image_frontal, image_sagittal = extract_mip(image)

    if mode == "sagittal":
        image = image_sagittal
    else:
        image = image_frontal

    return preprocess_mip_for_slice_detection(image, spacing, target_spacing)


def extract_mip(image, d=10, s=40):
    image_c = image.copy()

    image_c[
        :,
        :s,
    ] = 0
    image_c[
        :,
        -s:,
    ] = 0
    image_c[:, :, :s] = 0
    image_c[:, :, -s:] = 0

    (_, _, Z) = np.meshgrid(
        range(image.shape[1]), range(image.shape[0]), range(image.shape[2])
    )
    M = Z * (image_c > 0)
    M = M.sum(axis=2) / (image_c > 0).sum(axis=2)
    M[np.isnan(M)] = 0
    mask = M > 0
    c = int(np.mean(M[mask]))

    image_frontal = np.max(image_c, axis=1)

    image_sagittal = np.max(image_c[:, :, c - d : c + d], axis=2)[::-1, :]

    return image_frontal, image_sagittal


def run():
    parser = ArgumentParser()
    parser.add_argument("--dicom_dir")
    parser.add_argument("--target_spacing", type=float, default=0.5)
    args = parser.parse_args()

    reader = sitk.ImageSeriesReader()

    dicom_names = reader.GetGDCMSeriesFileNames(args.dicom_dir)
    reader.SetFileNames(dicom_names)

    image = reader.Execute()

    mip = preprocess_sitk_image_for_slice_detection(
        image, target_spacing=args.target_spacing
    )

    mip = np.dstack([mip[:, :, np.newaxis]] * 3).astype(np.uint8)
    imwrite("mip.jpg", mip)


if __name__ == "__main__":
    run()