Prepare_data.py

from __future__ import division
import nibabel as nib
import numpy as np
import Reza_functions as RF
import nibabel as nib
import glob
import os


# Define Train data and mask
Data_train   = []
Mask_train   = []
Maska_train  = []
FOV_train    = []

idx_count =1
Tr_add = '3d_images'

Tr_list = glob.glob(Tr_add+'/*.gz')

for idx in range(len(Tr_list)):
    b = Tr_list[idx]
    a = b[len(Tr_add)+1:len(Tr_add)+4]
    if a =='IMG':
        print(idx_count)
        a = b[len(Tr_add)+5:len(b)]
        add = (Tr_add+'/MASK_' + a) 
        vol = nib.load(Tr_list[idx])
        seg = nib.load(add)
        # Get the axials images and corresponding masks
        vol_ims, lung, around_lung, FOV = RF.return_axials(vol, seg)          
        segmentation  = seg.get_data()
        # Insert samples to the Train data, which has the segmentation label
        for idx in range(vol.shape[0]):
            if ~( np.sum(np.sum(np.sum(segmentation[idx, :, :]))) == 0):
                    
                Data_train.append(vol_ims [idx, :, :])
                Mask_train.append(lung[idx, :, :])
                Maska_train.append(around_lung[idx, :, :])               
                FOV_train.append(FOV[idx, :, :])               
        idx_count += 1
        
Data_train  = np.array(Data_train)
Mask_train  = np.array(Mask_train)
Maska_train = np.array(Maska_train)
FOV_train   = np.array(FOV_train)
print(2)
# We use 70% of the data for training and 30% for test
alpha = np.int16(np.floor(Data_train.shape[0]* 0.7))
en_d  = Data_train.shape[0]

Train_img      = Data_train[0:alpha,:,:]
Test_img       = Data_train[alpha:en_d,:,:]

Train_mask     = Mask_train[0:alpha,:,:]
Test_mask      = Mask_train[alpha:en_d,:,:]

Train_maska     = Maska_train[0:alpha,:,:]
Test_maska      = Maska_train[alpha:en_d,:,:]

FOV_tr     = FOV_train[0:alpha,:,:]
FOV_te      = FOV_train[alpha:en_d,:,:]
print(3)
folder = './processed_data/'
if not os.path.exists(folder):
    os.makedirs(folder)
print(4)    
np.save(folder+'data_train' , Train_img)
np.save(folder+'data_test'  , Test_img)
np.save(folder+'mask_train' , Train_mask)
np.save(folder+'mask_test'  , Test_mask)

np.save(folder+'Train_maska' , Train_maska)
np.save(folder+'Test_maska'  , Test_maska)
np.save(folder+'FOV_tr'      , FOV_tr)
np.save(folder+'FOV_te'      , FOV_te)