3_patchify_images.py

'''
Pre-process the data to extract patches
Input: A csv file containing path to input files

'''
import argparse
import os
import sys
import math
import numpy as np
import SimpleITK as sitk
import pandas as pd
import multiprocessing as mp

lowerThreshold = -1024
upperThreshold = 240

def convert_to_isotropic(inputVolume, isoSpacing=1.0):

    inputSpacing = inputVolume.GetSpacing()
    inputSize = inputVolume.GetSize()
    #Resample the images to make them iso-tropic
    resampleFilter = sitk.ResampleImageFilter()
    T = sitk.Transform()
    T.SetIdentity()
    resampleFilter.SetTransform(T)
    resampleFilter.SetInterpolator(sitk.sitkBSpline)
    resampleFilter.SetDefaultPixelValue(float(-1024))
    # isoSpacing = 1 #math.sqrt(inputSpacing[2] * inputSpacing[0])
    resampleFilter.SetOutputSpacing((isoSpacing,isoSpacing,isoSpacing))
    resampleFilter.SetOutputOrigin(inputVolume.GetOrigin())
    resampleFilter.SetOutputDirection(inputVolume.GetDirection())
    dx = int(inputSize[0] * inputSpacing[0] / isoSpacing)
    dy = int(inputSize[1] * inputSpacing[1] / isoSpacing)
    dz = int((inputSize[2] - 1 ) * inputSpacing[2] / isoSpacing)
    resampleFilter.SetSize((dx,dy,dz))
    try:
        resampleVolume = resampleFilter.Execute(inputVolume)
    except Error as err:
        print("Resample failed: " + str(imageFilePath) )
        print(err.decode(encoding='UTF-8'))
        return None

    return resampleVolume

def pad_img(input_img, image_lowest_intensity=-1024):
    lower_bound = [30] * 3
    upper_bound = [30] * 3
    cpf = sitk.ConstantPadImageFilter()
    cpf.SetConstant(image_lowest_intensity)
    cpf.SetPadLowerBound(lower_bound)
    cpf.SetPadUpperBound(upper_bound)
    input_img = cpf.Execute(input_img)
    return input_img

def Image2Patch(inputImg, step_size, patch_size, registered_patch_loc):
    """ This function converts image to patches. 
        Here is the input of the function:
          inputImg : input image. This should be simpleITK object
          patchSize : size of the patch. It should be array of three scalar
        Here is the output of the function:
          patchImgData : It is a list containing the patches of the image
          patchLblData : Is is a list containing the patches of the label image
          
    """
    patch_vol = patch_size[0]*patch_size[1]*patch_size[2]
    patch_img_data = []
    patch_loc = []

    for i in range(registered_patch_loc.shape[0]):
        x, y, z = registered_patch_loc[i].tolist()
        #print(x,y,z)
        patchImg = sitk.RegionOfInterest(inputImg, size=patch_size, index=[x,y,z])
        npLargePatchImg = sitk.GetArrayFromImage(patchImg)
        patch_img_data.append(npLargePatchImg.copy())
        patch_loc.append([x, y, z])
    
    patch_img_data = np.asarray(patch_img_data)
    patch_loc = np.asarray(patch_loc)
    return patch_img_data, patch_loc

def extract_patch(isoRawImage_file, altas_patch_loc):
    #Read the input isotropic image volume
    isoRawImage = sitk.ReadImage(isoRawImage_file)
    isoRawImage = convert_to_isotropic(isoRawImage)
    isoRawImage = pad_img(isoRawImage)
    npIsoRawImage = sitk.GetArrayFromImage(isoRawImage)
    #print(npIsoRawImage.shape)

    # Thresholding the isotropic raw image
    npIsoRawImage[npIsoRawImage > upperThreshold] = upperThreshold
    npIsoRawImage[npIsoRawImage < lowerThreshold] = lowerThreshold

    thresholdIsoRawImage = sitk.GetImageFromArray(npIsoRawImage)
    thresholdIsoRawImage.SetOrigin(isoRawImage.GetOrigin())
    thresholdIsoRawImage.SetSpacing(isoRawImage.GetSpacing())
    thresholdIsoRawImage.SetDirection(isoRawImage.GetDirection())
    
    # Prepare registered patch location
    registered_patch_loc = []
    affine_trans=sitk.ReadTransform("./INSP2Atlas/transform/"\
        + isoRawImage_file.split('/')[-1][:-7]+"_Reg_Atlas_Affine_0GenericAffine.mat")
    for i in range(altas_patch_loc.shape[0]):
        physical_cor_on_fixed = tuple(altas_patch_loc[i])
        physical_cor_on_moving = affine_trans.TransformPoint(physical_cor_on_fixed)
        index_on_moving = isoRawImage.TransformPhysicalPointToIndex(physical_cor_on_moving)
        registered_patch_loc.append(list(index_on_moving))
    registered_patch_loc = np.array(registered_patch_loc)

    #Extract Patches
    # Generate Patches of the masked Image
    patchImgData, patch_loc = Image2Patch(thresholdIsoRawImage, \
            [step_size]*3, [patch_size]*3, registered_patch_loc)
    return patchImgData, patch_loc

def prep_adjacency_matrix(patch_loc):
    adj = []
    for i in range(patch_loc.shape[0]):
        adj_row = np.zeros((patch_loc.shape[0],))
        dist = np.abs(patch_loc - patch_loc[i])
        max_side_dist = dist.max(1)
        dist = dist[max_side_dist<patch_size,:]
        volume = np.abs(dist-patch_size)
        volume = volume[:,0] * volume[:,1] * volume[:,2]
        #print(volume.shape)
        #print(adj_row[max_side_dist<patch_size].shape)
        adj_row[max_side_dist<patch_size] = volume / (patch_size**3)
        adj.append(adj_row.transpose())
    adj = np.asarray(adj)
    #adj = (adj / np.sum(adj, 0)).transpose()
    return adj

def run(start, end, batch_index): 
    df = pd.read_csv(input_csv)
    df = df[~df['image'].isnull()]

    # Prepare physical coord of patch location
    fixed_img = sitk.ReadImage(atlas_image)
    altas_patch_loc_temp=[]
    for i in range(altas_patch_loc.shape[0]):
        altas_patch_loc_temp.append(list(fixed_img.TransformIndexToPhysicalPoint(tuple(altas_patch_loc[i,:].tolist()))))
    altas_patch_loc=np.array(altas_patch_loc_temp)
    del altas_patch_loc_temp, fixed_img

    for i in range(start,end):
        row = df.iloc[i]   
        subject_id = row['sid']
        #if subject_id != '19676E':
        #    continue
        print(row['image'])
        output_basename = row['image'].split('/')[-1][:-24]
        isotropicFileName = "./INSP2Atlas/clamped/"+output_basename+".nii.gz"
        patchFileName = os.path.join(output_dir, 'patch', output_basename+'_patch.npy')

        if not os.path.exists(isotropicFileName):
            print(output_basename, "image not found")
            continue
        if not os.path.exists("./INSP2Atlas/transform/"\
            + output_basename +"_Reg_Atlas_Affine_0GenericAffine.mat"):
            print(output_basename, "mat not found")
            continue
        if os.path.getsize(isotropicFileName)/(1024.**2) < 16:
            continue

        if not os.path.exists(patchFileName):
            try:
                patchImgData, patch_loc = extract_patch(isotropicFileName, altas_patch_loc)

                adj = prep_adjacency_matrix(patch_loc)
                np.save(patchFileName, patchImgData)
                np.save(os.path.join(output_dir, 'patch_loc', output_basename+'_patch_loc.npy'), patch_loc)
                np.save(os.path.join(output_dir, 'adj', output_basename+'_adj.npy'), adj)
            except Exception as e:
                print(e)
                print("Failed in extract patch: " + str(output_basename) )
                continue
    return
                
def main(argv): 
    global output_dir
    global input_csv
    global patch_size 
    global step_size
    global atlas_image, atlas_patch_loc
    global lowerThreshold, upperThreshold

    #Parse the arguments
    parser = argparse.ArgumentParser(description='Data Preprocessing')
    parser.add_argument('--atlas_image', type=str, default="19676E_INSP_STD_JHU_COPD_BSpline_Iso1.0mm.nii.gz")
    parser.add_argument('--atlas_patch_loc', type=str, default="./patch_data_32_6_reg/19676E_INSP_STD_JHU_COPD_BSpline_Iso1_patch_loc.npy")
    parser.add_argument('--lowerThreshold', type=int, default=-1024)
    parser.add_argument('--upperThreshold', type=int, default=240)
    parser.add_argument('-i', '--input_csv', type=str,\
                        default='/pghbio/dbmi/batmanlab/Data/COPDGene/Database/Final_Status_Phase-1_18_11_2019_13_12_04.csv',\
                        help = 'Input csv with patient_id, isotropic volume path and lung segmentation path.')
    parser.add_argument('-o', '--output_dir', type=str, default='./patch_data_32_6_reg/',\
                        help='Directory where intermediate and final files are saved.')
    parser.add_argument('-p', '--num_processor', type=int, default=16,\
                        help='Preprocess files in parallel. 1: no-parallel, n: number of nodes for parallel use')
    parser.add_argument('-s', '--patch_size', type=int, default=32, help='The size of the 3D patch.')
    parser.add_argument('-l', '--step_size', type=int, default=26, help='The overlap between consecutive patches.')
    
    
    args = parser.parse_args()
    input_csv = args.input_csv
    output_dir = args.output_dir
    num_processor = args.num_processor
    patch_size = args.patch_size
    step_size = args.step_size
    atlas_image = args.atlas_image
    atlas_patch_loc = args.atlas_patch_loc
    lowerThreshold = args.lowerThreshold
    upperThreshold = args.upperThreshold

    df = pd.read_csv(input_csv)
    df = df[~df['image'].isnull()]

    start = 0
    end = df.shape[0] 

    if num_processor > 1:
        processes = []
        count = math.ceil((end - start)/num_processor)
        for i in range(num_processor):
            if start + count > end:
                count = end - start
            processes.append(mp.Process(target=run, args=(start, start + count, args.batch_index)))
            start += count
        for p in processes:
            p.start()
    else:
        run(start, end)

if __name__ == '__main__':
    main(sys.argv[1:])