# -*- coding: utf-8 -*-
"""
Created on Fri Oct 14 01:47:13 2016
@author: seeker105
"""
import os.path
import random
import pylab
import numpy as np
from glob import glob
import SimpleITK as sitk
from Nyul import IntensityRangeStandardization
import sys
import timeit
from sklearn.feature_extraction.image import extract_patches_2d
from skimage import color
class Pipeline(object):
''' The Pipeline for loading images for all patients and all modalities
1)find_training_patches: finds the training patches for a particular class
INPUT:
1) The filepath 'path': Directory of the image database. It contains the slices of training image slices
'''
def __init__(self, path_train = '', path_test = '' , mx_train = 1000000, mx_tst = 1000000):
self.path_train = path_train
self.path_test = path_test
self.scans_train, self.scans_test, self.train_im, self.test_im = self.read_scans(mx_train, mx_tst)
def read_scans(self, mx_train, mx_test):
scans_train = glob(self.path_train + r'/*.mha')
scans_test = glob(self.path_test + r'/*.mha')
train_im = [sitk.GetArrayFromImage(sitk.ReadImage(scans_train[i])) for i in xrange(min(len(scans_train), mx_train))]
test_im = [sitk.GetArrayFromImage(sitk.ReadImage(scans_test[i])) for i in xrange(min(len(scans_test), mx_test))]
return scans_train, scans_test, np.array(train_im), np.array(test_im)
def n4itk(self, img):
img = sitk.Cast(img, sitk.sitkFloat32)
img_mask = sitk.BinaryNot(sitk.BinaryThreshold(img, 0, 0)) ## Create a mask spanning the part containing the brain, as we want to apply the filter to the brain image
corrected_img = sitk.N4BiasFieldCorrection(img, img_mask)
return corrected_img
'''def find_all_train(self, classes, d = 4, h = 33, w = 33):
mn = 300000000000000
#load all labels
im_ar = []
for i in self.pathnames_train:
im_ar.append([sitk.GetArrayFromImage(sitk.ReadImage(idx)) for idx in i])
im_ar = np.array(im_ar)
for i in xrange(classes):
mn = min(mn, len(np.argwhere(im_ar[i]==i)))
'''
def sample_training_patches(self, num_patches, class_nm, d = 4, h = 33, w = 33):
''' Creates the input patches and their labels for training CNN. The patches are 4x33x33
and the label for a patch equals to the label for the central pixel of the patch.
INPUT:
1) num_patches: The number of patches required of the class.
2) class_nm: The index of class label for which we are finding patches.
3) d, h, w: number of channels, height and width of patch
OUTPUT:
1) patches: The list of all patches of dimensions d, h, w.
2) labels: The list of labels for each patch. Label for a patch corresponds to the label
of the central pixel of that patch.
'''
#find patches for training
patches, labels = [], np.full(num_patches, fill_value = class_nm, dtype = np.int32)
count = 0
# convert gt_im to 1D and save shape
gt_im = np.swapaxes(self.train_im, 0, 1)[4] #swap axes to make axis 0 represent the modality and axis 1 represent the slice. take the ground truth
#take flair image as mask
msk = np.swapaxes(self.train_im, 0, 1)[0]
tmp_shp = gt_im.shape
gt_im = gt_im.reshape(-1)
msk = msk.reshape(-1)
# maintain list of 1D indices where label = class_nm
indices = np.squeeze(np.argwhere((gt_im == class_nm) & (msk != 0.)))
# shuffle the list of indices of the class
st = timeit.default_timer()
np.random.shuffle(indices)
print 'shuffling of label {} took :'.format(class_nm), timeit.default_timer()-st
#reshape gt_im
gt_im = gt_im.reshape(tmp_shp)
st = timeit.default_timer()
#find the patches from the images
i = 0
pix = len(indices)
while (count<num_patches) and (pix>i):
#print (count, ' cl:' ,class_nm)
#sys.stdout.flush()
#randomly choose an index
ind = indices[i]
i+= 1
#reshape ind to 3D index
ind = np.unravel_index(ind, tmp_shp)
#print ind
#sys.stdout.flush()
#find the slice index to choose from
slice_idx = ind[0]
#load the slice from the label
l = gt_im[slice_idx]
# the centre pixel and its coordinates
p = ind[1:]
#construct the patch by defining the coordinates
p_x = (p[0] - h/2, p[0] + (h+1)/2)
p_y = (p[1] - w/2, p[1] + (w+1)/2)
#check if the pixels are in range
if p_x[0]<0 or p_x[1]>l.shape[0] or p_y[0]<0 or p_y[1]>l.shape[1]:
continue
#take patches from all modalities and group them together
tmp = self.train_im[slice_idx][0:4, p_x[0]:p_x[1], p_y[0]:p_y[1]]
patches.append(tmp)
count+=1
print 'finding patches of label {} took :'.format(class_nm), timeit.default_timer()-st
patches = np.array(patches)
return patches, labels
def training_patches(self, num_patches, classes = 5, d = 4, h = 33, w = 33):
'''Creates the input patches and their labels for training CNN. The patches are 4x33x33
and the label for a patch corresponds to the label for the central voxel of the patch. The
data will be balanced, with the number of patches being the same for each class
INPUT:
1) classes: number of all classes in the segmentation
2) num_patches: number of patches for each class
3) d, h, w : channels, height and width of the patches
OUTPUT:
1) all_patches: numpy array of all class patches of the shape 4x33x33
2) all_labels : numpy array of the all_patches labels
'''
patches, labels, mu, sigma = [], [], [], []
for idx in xrange(classes):
p, l = self.sample_training_patches(num_patches[idx], idx, d, h, w)
patches.append(p)
labels.append(l)
patches = np.vstack(np.array(patches))
patches_by_channel = np.swapaxes(patches, 0, 1)
for seq, i in zip(patches_by_channel, xrange(d)):
avg = np.mean(seq)
std = np.std(seq)
patches_by_channel[i] = (patches_by_channel[i] - avg)/std
mu.append(avg)
sigma.append(std)
patches = np.swapaxes(patches_by_channel, 0, 1)
return patches, np.array(labels).reshape(-1), np.array(mu), np.array(sigma)
def test_patches(img , mu, sigma, d = 4, h = 33, w = 33):
''' Creates patches of image. Returns a numpy array of dimension number_of_patches x d.
INPUT:
1)img: a 3D array containing the all modalities of a 2D image.
2)d, h, w: see above
OUTPUT:
tst_arr: ndarray of all patches of all modalities. Of the form number of patches x modalities
'''
#list of patches
p = []
msk = (img[0]+img[1]+img[2]+img[3])!=0. #mask using FLAIR channel
msk = msk[16:-16, 16:-16] #crop the mask to conform to the rebuilt image after prediction
msk = msk.reshape(-1)
for i in xrange(len(img)):
plist = extract_patches_2d(img[i], (h, w))
plist = (plist - mu[i])/sigma[i]
p.append(plist[msk]) #only take patches with brain mask!=0
return np.array(p).swapaxes(0,1)
def reconstruct_labels(msk, pred_list):
im = np.full((208, 208), 0.)
msk = msk[16:-16, 16:-16]
im[msk] = np.array(pred_list)
im = np.pad(im, (16, 16), mode='edge')
return im
Datasets
Models
