convert_dicom_to_npy.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 qiang.zhou qiang.zhou@Macbook
#
# Distributed under terms of the MIT license.

import numpy as np
import os
import pydicom
import cv2

""" Return a int32 dtype numpy image of a directory. """
def load_16bit_dicom_images(path, verbose=True):
    slices = [pydicom.read_file(os.path.join(path, s), force=True)
                        for s in os.listdir(path)]
    slices.sort(key=lambda x: float(x.ImagePositionPatient[2]))

    #import pdb
    #pdb.set_trace()
    total_cnt = len(slices)

    slice_thickness = np.abs(slices[0].ImagePositionPatient[2]-
                            slices[1].ImagePositionPatient[2])
    # Remove duplicated slices
    ignore_cnt = 0
    if slice_thickness == 0:
        unique_slices = [slices[0]]
        start_pos = slices[0].ImagePositionPatient[2]

        for s in slices[1:]:
            if s.ImagePositionPatient[2] != start_pos:
                unique_slices.append(s)
                start_pos = s.ImagePositionPatient[2]
            else:
                ignore_cnt += 1
        slices = unique_slices
        slice_thickness = np.abs(slices[0].ImagePositionPatient[2]-
                                slices[1].ImagePositionPatient[2])

    if verbose:
        print ("[INFO] Total/Ignore/Reserved {}/{}/{} CTAs, slice_thickess {}...\
              ".format(total_cnt, ignore_cnt, len(slices), slice_thickness))

    for s in slices:
        s.SliceThickness = slice_thickness

    # Fix HU
    image = np.stack([s.pixel_array for s in slices])
    image = image.astype(np.int16)
    image[image <= -2000] = 0
    image[image >= 3000] = 3000 

    for slice_number in range(len(slices)):
        intercept = slices[slice_number].RescaleIntercept   # -1024
        slope = slices[slice_number].RescaleSlope           # 1

        if slope != 1:
            image[slice_number] = slope * image[slice_number].astype(np.float64)
            image[slice_number] = image[slice_number].astype(np.int16)
        image[slice_number] += np.int16(intercept)

    numpySpacing = [float(slice_thickness)] + [float(x) for x in slices[0].PixelSpacing]

    numpyImage = np.array(image, dtype=np.int16)            # SxHxW
    numpySpacing = np.array(numpySpacing, dtype=np.float)

    return numpyImage, numpySpacing

def normalize_16bit_dicom_images(cta_image, HU_window=np.array([-1000., 400.]), bound_values=[0, 1]):
    # Outlier
    #mid = (0-HU_window[0])/(HU_window[1] - HU_window[0])
    #cta_image[cta_image == 0] = HU_window[0]

    th_cta_image = (cta_image-HU_window[0])/(HU_window[1] - HU_window[0])
    th_cta_image[th_cta_image < 0] = bound_values[0]
    th_cta_image[th_cta_image >= 1] = bound_values[1]
    th_cta_image_mask = (th_cta_image*255).astype('uint8')
    return th_cta_image_mask


def resample(image, spacing, new_spacing=[1,1,1]):
    resize_factor = spacing / new_spacing
    new_real_shape = image.shape * resize_factor
    new_shape = np.round(new_real_shape)
    real_resize_factor = new_shape / image.shape
    new_spacing = spacing / real_resize_factor

    image = scipy.ndimage.interpolation.zoom(image, 
                                            real_resize_factor,
                                            mode="nearest")
    return image

isotropic_resolution = np.array([1,1,1])
box3d_margin = 5            # pixels
clip_ratio = 0.01
HU_window = np.array([-1200, 600])

# Lambda funcs
listdironly = lambda d: [os.path.join(d, o) for o in os.listdir(d) if os.path.isdir(os.path.join(d, o))]
dumpimgs = lambda x, p: [cv2.imwrite("debug/{}_{:05d}.png".format(p, i), xi) for i, xi in enumerate(x)]

# 加一个圆形的mask是因为CT机器一般是圆形扫描的，会导致圆外部的HU值很诡异, 
#这个圆应该比他标准CT的小一些，防止边缘漏掉了
from skimage.draw import circle
circle_mask = np.zeros((512, 512), dtype=np.uint8) 
rr, cc = circle(255, 255, 253)  # the standard is 256, we use 253
circle_mask[rr, cc] = 1
circle_mask_inv = 1-circle_mask

# 最标准的结果，没有任何额外处理
# 1. 把肺部分割出来，失败也不管
# 2. 归一化到1mmx1mmx1mm

def main(path):

    unique_id = path.split('/')[-1]

    print (f"[INFO] Processing {path} - {unique_id}")
    sliceim_pixels, spacing = load_16bit_dicom_images(path)

    # Use a circle to filter out unrelated area
    for i in range(len(sliceim_pixels)):
        sliceim_pixels[i] += -2000 * circle_mask_inv

    # Resample to 1x1x1 and compute the enclosed box
    sliceim_isotropic = resample(sliceim_pixels, spacing, isotropic_resolution)
    print (f"\t[SUBINFO] lung images' shape: {sliceim_isotropic.shape}")

    # Thresholding specified HU areas and obtained 3D box
    sliceim_mask = normalize_16bit_dicom_images(sliceim_isotropic, HU_window=HU_window)
    np.save(os.path.join(f"{unique_id}.npy"), sliceim_mask)

if __name__ == "__main__":
    assert False, "Please modify the script configurations..."
    CTA_dirpath = "YOUR_DICOM_DIRPATH"
    main(CTA_dirpath)