├── __init__.py ├── utils.py ├── metric.py ├── average_ensembles.py ├── submission.py ├── README.md ├── current.py ├── data.py ├── keras_plus.py ├── current_ensemble.py ├── train_kfold.py ├── u_model.py ├── augmentation.py ├── train.py └── train_generator.py /__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import cPickle as pickle 2 | 3 | def load_pickle(file_path): 4 | data = None 5 | with open (file_path,"rb") as dumpFile: 6 | data = pickle.load(dumpFile) 7 | return data 8 | 9 | def save_pickle(file_path, data): 10 | with open (file_path,"wb") as dumpFile: 11 | pickle.dump(data, dumpFile, pickle.HIGHEST_PROTOCOL) 12 | 13 | def count_enum(words): 14 | wdict = {} 15 | get = wdict.get 16 | for word in words: 17 | wdict[word] = get(word, 0) + 1 18 | return wdict 19 | -------------------------------------------------------------------------------- /metric.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import numpy as np 3 | from keras import backend as K 4 | smooth = 1 5 | 6 | 7 | def mean_length_error(y_true, y_pred): 8 | y_true_f = K.sum(K.round(K.flatten(y_true))) 9 | y_pred_f = K.sum(K.round(K.flatten(y_pred))) 10 | delta = (y_pred_f - y_true_f) 11 | return K.mean(K.tanh(delta)) 12 | 13 | def dice_coef(y_true, y_pred): 14 | y_true_f = K.flatten(y_true) 15 | y_pred_f = K.flatten(y_pred) 16 | intersection = K.sum(y_true_f * y_pred_f) 17 | return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth) 18 | 19 | def dice_coef_loss(y_true, y_pred): 20 | return -dice_coef(y_true, y_pred) 21 | 22 | def np_dice_coef(y_true, y_pred): 23 | tr = y_true.flatten() 24 | pr = y_pred.flatten() 25 | return (2. * np.sum(tr * pr) + smooth) / (np.sum(tr) + np.sum(pr) + smooth) 26 | 27 | 28 | def main(): 29 | a = np.random.random((420,100)) 30 | b = np.random.random((420,100)) 31 | # print a.flatten().shape 32 | res = np_dice_coef(a,b ) 33 | print res 34 | 35 | 36 | if __name__ == '__main__': 37 | sys.exit(main()) 38 | -------------------------------------------------------------------------------- /average_ensembles.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sys 3 | from u_model import IMG_COLS as img_cols, IMG_ROWS as img_rows 4 | from train import Learner 5 | 6 | ensembles = { 7 | 'ens2': (8, 'best/ens2/res3/'), 8 | 'ens3': (6, 'best/ens3/res3/'), 9 | 'ens4': (6, 'best/ens4/res3/'), 10 | 'ens5': (8, 'best/ens5/res3/'), 11 | 'ens7': (6, 'best/ens7/res3/'), 12 | 'ens8': (5, 'best/ens8/res3/'), 13 | } 14 | 15 | 16 | def main(): 17 | kfold_masks, kfold_prob = [], [] 18 | weigths = [] 19 | for name, (kfold, prefix) in ensembles.iteritems(): 20 | print 'Loading name=%s, prefix=%s, kfold=%d' % (name, prefix, kfold) 21 | ens_x_mask = np.load(prefix + 'imgs_mask_test.npy') 22 | ens_x_prob = np.load(prefix + 'imgs_mask_exist_test.npy') 23 | kfold_masks.append(ens_x_mask) 24 | kfold_prob.append(ens_x_prob) 25 | weigths.append(kfold) 26 | # 27 | total_weight = float(sum(weigths)) 28 | total_cnt = len(weigths) 29 | dlen = len(kfold_masks[0]) 30 | res_masks = np.ndarray((dlen, 1, img_rows, img_cols), dtype=np.float32) 31 | res_probs = np.ndarray((dlen, ), dtype=np.float32) 32 | 33 | for i in xrange(dlen): 34 | masks = np.ndarray((total_cnt, 1, img_rows, img_cols), dtype=np.float32) 35 | probs = np.ndarray((total_cnt, ), dtype=np.float32) 36 | for k in xrange(total_cnt): 37 | masks[k] = weigths[k] * kfold_masks[k][i] 38 | probs[k] = weigths[k] * kfold_prob[k][i] 39 | res_masks[i] = np.sum(masks, 0)/total_weight 40 | res_probs[i] = np.sum(probs)/total_weight 41 | print 'Saving', Learner.test_mask_res, Learner.test_mask_exist_res 42 | np.save(Learner.test_mask_res, res_masks) 43 | np.save(Learner.test_mask_exist_res, res_probs) 44 | 45 | 46 | if __name__ == '__main__': 47 | sys.exit(main()) 48 | 49 | -------------------------------------------------------------------------------- /submission.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import sys, os 3 | import numpy as np 4 | import cv2 5 | from data import image_cols, image_rows, load_test_ids 6 | from train import Learner 7 | 8 | 9 | def prep(img): 10 | img = img.astype('float32') 11 | img = cv2.resize(img, (image_cols, image_rows)) 12 | img = cv2.threshold(img, 0.5, 1., cv2.THRESH_BINARY)[1].astype(np.uint8) 13 | return img 14 | 15 | def run_length_enc(label): 16 | from itertools import chain 17 | x = label.transpose().flatten() 18 | y = np.where(x > 0)[0] 19 | if len(y) < 10: # consider as empty 20 | return '' 21 | z = np.where(np.diff(y) > 1)[0] 22 | start = np.insert(y[z+1], 0, y[0]) 23 | end = np.append(y[z], y[-1]) 24 | length = end - start 25 | res = [[s+1, l+1] for s, l in zip(list(start), list(length))] 26 | res = list(chain.from_iterable(res)) 27 | return ' '.join([str(r) for r in res]) 28 | 29 | 30 | def submission(): 31 | imgs_id_test = load_test_ids() 32 | 33 | print ('Loading test_mask_res from %s' % Learner.test_mask_res) 34 | imgs_test = np.load(Learner.test_mask_res) 35 | print ('Loading imgs_exist_test from %s' % Learner.test_mask_exist_res) 36 | imgs_exist_test = np.load(Learner.test_mask_exist_res) 37 | 38 | argsort = np.argsort(imgs_id_test) 39 | imgs_id_test = imgs_id_test[argsort] 40 | imgs_test = imgs_test[argsort] 41 | imgs_exist_test = imgs_exist_test[argsort] 42 | 43 | total = imgs_test.shape[0] 44 | ids = [] 45 | rles = [] 46 | for i in xrange(total): 47 | img = imgs_test[i, 0] 48 | img_exist = imgs_exist_test[i] 49 | img = prep(img) 50 | new_prob = (img_exist + min(1, np.sum(img)/10000.0 )* 5 / 3)/2 51 | if np.sum(img) > 0 and new_prob < 0.5: 52 | img = np.zeros((image_rows, image_cols)) 53 | 54 | rle = run_length_enc(img) 55 | 56 | rles.append(rle) 57 | ids.append(imgs_id_test[i]) 58 | 59 | if i % 1000 == 0: 60 | print('{}/{}'.format(i, total)) 61 | 62 | file_name = os.path.join(Learner.res_dir, 'submission.csv') 63 | 64 | with open(file_name, 'w+') as f: 65 | f.write('img,pixels\n') 66 | for i in xrange(total): 67 | s = str(ids[i]) + ',' + rles[i] 68 | f.write(s + '\n') 69 | 70 | def main(): 71 | submission() 72 | 73 | 74 | if __name__ == '__main__': 75 | sys.exit(main()) 76 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Ultrasound nerve segmentation using Keras (1.0.7) 2 | Kaggle Ultrasound Nerve Segmentation competition [Keras] 3 | 4 | #Install (Ubuntu {14,16}, GPU) 5 | 6 | cuDNN required. 7 | 8 | ###Theano 9 | - http://deeplearning.net/software/theano/install_ubuntu.html#install-ubuntu 10 | - sudo pip install pydot-ng 11 | 12 | In ~/.theanorc 13 | ``` 14 | [global] 15 | device = gpu0 16 | [dnn] 17 | enabled = True 18 | ``` 19 | 20 | ###Keras 21 | - sudo apt-get install libhdf5-dev 22 | - sudo pip install h5py 23 | - sudo pip install pydot 24 | - sudo pip install nose_parameterized 25 | - sudo pip install keras 26 | 27 | In ~/.keras/keras.json (it's very important, the project was running on theano backend, and some issues are possible in case of TensorFlow) 28 | ``` 29 | { 30 | "image_dim_ordering": "th", 31 | "epsilon": 1e-07, 32 | "floatx": "float32", 33 | "backend": "theano" 34 | } 35 | ``` 36 | 37 | ###Python deps 38 | - sudo apt-get install python-opencv 39 | - sudo apt-get install python-sklearn 40 | 41 | #Prepare 42 | 43 | Place train and test data into '../train' and '../test' folders accordingly. 44 | 45 | ``` 46 | mkdir np_data 47 | python data.py 48 | ``` 49 | 50 | #Training 51 | 52 | Single model training. 53 | ``` 54 | python train.py 55 | ``` 56 | Results will be generatated in "res/" folder. res/unet.hdf5 - best model 57 | 58 | Generate submission: 59 | ``` 60 | python submission.py 61 | ``` 62 | 63 | Generate predection with a model in res/unet.hdf5 64 | ``` 65 | python current.py 66 | ``` 67 | 68 | #Model 69 | 70 | Motivation's explained in my internal pres (slides: http://www.slideshare.net/Eduardyantov/ultrasound-segmentation-kaggle-review) 71 | 72 | I used U-net like architecture (http://arxiv.org/abs/1505.04597). Main differences: 73 | - inception blocks instead of VGG like 74 | - Conv with stride instead of MaxPooling 75 | - Dropout, p=0.5 76 | - skip connections from encoder to decoder layers with residual blocks 77 | - BatchNorm everywhere 78 | - 2 heads training: auxiliary branch for scoring nerve presence (in the middle of the network), one branch for segmentation 79 | - ELU activation 80 | - sigmoid activation in output 81 | - Adam optimizer, without weight regularization in layers 82 | - Dice coeff loss, average per batch, without smoothing 83 | - output layers - sigmoid activation 84 | - batch_size=64,128 (for GeForce 1080 and Titan X respectively) 85 | 86 | Augmentation: 87 | - flip x,y 88 | - random zoom 89 | - random channel shift 90 | - elastic transormation didn't help in this configuration 91 | 92 | Augmentation generator (generate augmented data on the fly for each epoch) didn't improve the score. 93 | For prediction augmented images were used. 94 | 95 | Validation: 96 | 97 | For some reason validation split by patient (which is proper in this competition) didn't work for me, probably due to bug in the code. So I used random split. 98 | 99 | Final prediction uses probability of a nerve presence: p_nerve = (p_score + p_segment)/2, where p_segment based on number of output pixels in the mask. 100 | 101 | #Results and technical aspects 102 | - On GPU Titan X an epoch took about 6 minutes. Training early stops at 15-30 epochs. 103 | - For batch_size=64 6Gb GPU memory is required. 104 | - Best single model achieved 0.694 LB score. 105 | - An ensemble of 6 different k-fold ensembles (k=5,6,8) scored 0.70399 106 | 107 | #Credits 108 | This code was originally based on https://github.com/jocicmarko/ultrasound-nerve-segmentation/ 109 | -------------------------------------------------------------------------------- /current.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sys 3 | from data import load_test_data 4 | from u_model import get_unet 5 | from keras.optimizers import Adam 6 | from train import preprocess, Learner 7 | #aug 8 | from u_model import IMG_COLS as img_cols, IMG_ROWS as img_rows 9 | from keras.preprocessing.image import flip_axis, random_channel_shift 10 | 11 | curry = lambda func, *args, **kw:\ 12 | lambda *p, **n:\ 13 | func(*args + p, **dict(kw.items() + n.items())) 14 | from keras.preprocessing.image import apply_transform, transform_matrix_offset_center 15 | 16 | 17 | def zoom(x, zoom_range, row_index=1, col_index=2, channel_index=0, 18 | fill_mode='nearest', cval=0.): 19 | zx, zy = zoom_range 20 | zoom_matrix = np.array([[zx, 0, 0], 21 | [0, zy, 0], 22 | [0, 0, 1]]) 23 | 24 | h, w = x.shape[row_index], x.shape[col_index] 25 | transform_matrix = transform_matrix_offset_center(zoom_matrix, h, w) 26 | x = apply_transform(x, transform_matrix, channel_index, fill_mode, cval) 27 | return x 28 | 29 | 30 | transforms = ( 31 | {'do': curry(flip_axis, axis=1), 'undo': curry(flip_axis, axis=1)}, 32 | {'do': curry(flip_axis, axis=2), 'undo': curry(flip_axis, axis=2)}, 33 | {'do': curry(zoom, zoom_range=(1.05, 1.05)), 'undo': curry(zoom, zoom_range=(1/1.05, 1/1.05))}, 34 | {'do': curry(zoom, zoom_range=(0.95, 0.95)), 'undo': curry(zoom, zoom_range=(1/0.95, 1/0.95))}, 35 | {'do': curry(random_channel_shift, intensity=5), 'undo': lambda x: x}, 36 | ) 37 | 38 | 39 | def run_test(): 40 | BS = 128 41 | print('Loading and preprocessing test data...') 42 | mean, std = Learner.load_meanstd() 43 | 44 | imgs_test = load_test_data() 45 | imgs_test = preprocess(imgs_test) 46 | 47 | imgs_test = imgs_test.astype('float32') 48 | imgs_test -= mean 49 | imgs_test /= std 50 | 51 | print('Loading saved weights...') 52 | model = get_unet(Adam(0.001)) 53 | print ('Loading weights from %s' % Learner.best_weight_path) 54 | model.load_weights(Learner.best_weight_path) 55 | 56 | print ('Augment') 57 | alen, dlen = len(transforms), len(imgs_test) 58 | test_x = np.ndarray((alen, dlen, 1, img_rows, img_cols), dtype=np.float32) 59 | for i in xrange(dlen): 60 | for j, transform in enumerate(transforms): 61 | test_x[j,i] = transform['do'](imgs_test[i].copy()) 62 | # 63 | print('Predicting masks on test data...') 64 | outputs = [] 65 | asis_res = model.predict(imgs_test, batch_size=BS, verbose=1) 66 | outputs.append(asis_res) 67 | for j, transform in enumerate(transforms): 68 | t_y = model.predict(test_x[j], batch_size=BS, verbose=1) 69 | outputs.append(t_y) 70 | # 71 | print('Analyzing') 72 | test_masks = np.ndarray((dlen, 1, img_rows, img_cols), dtype=np.float32) 73 | test_probs = np.ndarray((dlen, ), dtype=np.float32) 74 | for i in xrange(dlen): 75 | masks = np.ndarray((alen+1, 1, img_rows, img_cols), dtype=np.float32) 76 | probs = np.ndarray((alen+1, ), dtype=np.float32) 77 | for j, t_y in enumerate(outputs): 78 | mask, prob = t_y[0][i], t_y[1][i] 79 | if j: 80 | mask = transforms[j-1]['undo'](mask) 81 | masks[j] = mask 82 | probs[j] = prob 83 | # 84 | test_masks[i] = np.mean(masks, 0) 85 | test_probs[i] = np.mean(probs) 86 | 87 | print('Saving ') 88 | np.save(Learner.test_mask_res, test_masks) 89 | np.save(Learner.test_mask_exist_res, test_probs) 90 | 91 | def main(): 92 | run_test() 93 | 94 | if __name__ == '__main__': 95 | sys.exit(main()) 96 | -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | import os, sys 3 | import numpy as np 4 | import cv2 5 | 6 | image_rows = 420 7 | image_cols = 580 8 | 9 | _dir = os.path.join(os.path.realpath(os.path.dirname(__file__)), '') 10 | data_path = os.path.join(_dir, '../') 11 | preprocess_path = os.path.join(_dir, 'np_data') 12 | img_train_path = os.path.join(preprocess_path, 'imgs_train.npy') 13 | img_train_mask_path = os.path.join(preprocess_path, 'imgs_mask_train.npy') 14 | img_train_patients = os.path.join(preprocess_path, 'imgs_patient.npy') 15 | img_test_path = os.path.join(preprocess_path, 'imgs_test.npy') 16 | img_test_id_path = os.path.join(preprocess_path, 'imgs_id_test.npy') 17 | 18 | 19 | 20 | def load_test_data(): 21 | print ('Loading test data from %s' % img_test_path) 22 | imgs_test = np.load(img_test_path) 23 | return imgs_test 24 | 25 | def load_test_ids(): 26 | print ('Loading test ids from %s' % img_test_id_path) 27 | imgs_id = np.load(img_test_id_path) 28 | return imgs_id 29 | 30 | def load_train_data(): 31 | print ('Loading train data from %s and %s' % (img_train_path, img_train_mask_path)) 32 | imgs_train = np.load(img_train_path) 33 | imgs_mask_train = np.load(img_train_mask_path) 34 | return imgs_train, imgs_mask_train 35 | 36 | def load_patient_num(): 37 | print ('Loading patient numbers from %s' % img_train_patients) 38 | return np.load(img_train_patients) 39 | 40 | def get_patient_nums(string): 41 | pat, photo = string.split('_') 42 | photo = photo.split('.')[0] 43 | return int(pat), int(photo) 44 | 45 | def create_train_data(): 46 | train_data_path = os.path.join(data_path, 'train') 47 | images = filter((lambda image: 'mask' not in image), os.listdir(train_data_path)) 48 | total = len(images) 49 | 50 | imgs = np.ndarray((total, 1, image_rows, image_cols), dtype=np.uint8) 51 | imgs_mask = np.ndarray((total, 1, image_rows, image_cols), dtype=np.uint8) 52 | i = 0 53 | print('Creating training images...') 54 | img_patients = np.ndarray((total,), dtype=np.uint8) 55 | for image_name in images: 56 | if 'mask' in image_name: 57 | continue 58 | image_mask_name = image_name.split('.')[0] + '_mask.tif' 59 | patient_num = image_name.split('_')[0] 60 | img = cv2.imread(os.path.join(train_data_path, image_name), cv2.IMREAD_GRAYSCALE) 61 | img_mask = cv2.imread(os.path.join(train_data_path, image_mask_name), cv2.IMREAD_GRAYSCALE) 62 | 63 | imgs[i, 0] = img 64 | imgs_mask[i, 0] = img_mask 65 | img_patients[i] = patient_num 66 | if i % 100 == 0: 67 | print('Done: {0}/{1} images'.format(i, total)) 68 | i += 1 69 | print('Loading done.') 70 | np.save(img_train_patients, img_patients) 71 | np.save(img_train_path, imgs) 72 | np.save(img_train_mask_path, imgs_mask) 73 | print('Saving to .npy files done.') 74 | 75 | 76 | def create_test_data(): 77 | train_data_path = os.path.join(data_path, 'test') 78 | images = os.listdir(train_data_path) 79 | total = len(images) 80 | 81 | imgs = np.ndarray((total, 1, image_rows, image_cols), dtype=np.uint8) 82 | imgs_id = np.ndarray((total, ), dtype=np.int32) 83 | 84 | i = 0 85 | print('Creating test images...') 86 | for image_name in images: 87 | img_id = int(image_name.split('.')[0]) 88 | img = cv2.imread(os.path.join(train_data_path, image_name), cv2.IMREAD_GRAYSCALE) 89 | 90 | imgs[i, 0] = img 91 | imgs_id[i] = img_id 92 | 93 | if i % 100 == 0: 94 | print('Done: {0}/{1} images'.format(i, total)) 95 | i += 1 96 | print('Loading done.') 97 | 98 | np.save(img_test_path, imgs) 99 | np.save(img_test_id_path, imgs_id) 100 | print('Saving to .npy files done.') 101 | 102 | 103 | def main(): 104 | create_train_data() 105 | create_test_data() 106 | 107 | if __name__ == '__main__': 108 | sys.exit(main()) 109 | -------------------------------------------------------------------------------- /keras_plus.py: -------------------------------------------------------------------------------- 1 | from keras.callbacks import Callback 2 | from keras.callbacks import warnings 3 | import sys 4 | import numpy as np 5 | from keras import backend as K 6 | 7 | 8 | class AdvancedLearnignRateScheduler(Callback): 9 | ''' 10 | # Arguments 11 | monitor: quantity to be monitored. 12 | patience: number of epochs with no improvement 13 | after which training will be stopped. 14 | verbose: verbosity mode. 15 | mode: one of {auto, min, max}. In 'min' mode, 16 | training will stop when the quantity 17 | monitored has stopped decreasing; in 'max' 18 | mode it will stop when the quantity 19 | monitored has stopped increasing. 20 | ''' 21 | def __init__(self, monitor='val_loss', patience=0, 22 | verbose=0, mode='auto', decayRatio=0.5): 23 | super(Callback, self).__init__() 24 | 25 | self.monitor = monitor 26 | self.patience = patience 27 | self.verbose = verbose 28 | self.wait = 0 29 | self.decayRatio = decayRatio 30 | 31 | if mode not in ['auto', 'min', 'max']: 32 | warnings.warn('Mode %s is unknown, ' 33 | 'fallback to auto mode.' 34 | % (self.mode), RuntimeWarning) 35 | mode = 'auto' 36 | 37 | if mode == 'min': 38 | self.monitor_op = np.less 39 | self.best = np.Inf 40 | elif mode == 'max': 41 | self.monitor_op = np.greater 42 | self.best = -np.Inf 43 | else: 44 | if 'acc' in self.monitor: 45 | self.monitor_op = np.greater 46 | self.best = -np.Inf 47 | else: 48 | self.monitor_op = np.less 49 | self.best = np.Inf 50 | 51 | def on_epoch_end(self, epoch, logs={}): 52 | current = logs.get(self.monitor) 53 | 54 | current_lr = K.get_value(self.model.optimizer.lr) 55 | print(" \nLearning rate:", current_lr) 56 | if current is None: 57 | warnings.warn('AdvancedLearnignRateScheduler' 58 | ' requires %s available!' % 59 | (self.monitor), RuntimeWarning) 60 | 61 | if self.monitor_op(current, self.best): 62 | self.best = current 63 | self.wait = 0 64 | else: 65 | if self.wait >= self.patience: 66 | assert hasattr(self.model.optimizer, 'lr'), \ 67 | 'Optimizer must have a "lr" attribute.' 68 | current_lr = K.get_value(self.model.optimizer.lr) 69 | new_lr = current_lr * self.decayRatio 70 | if self.verbose > 0: 71 | print(' \nEpoch %05d: reducing learning rate' % (epoch)) 72 | sys.stderr.write(' \nnew lr: %.5f\n' % new_lr) 73 | K.set_value(self.model.optimizer.lr, new_lr) 74 | self.wait = 0 75 | 76 | self.wait += 1 77 | 78 | 79 | class LearningRateDecay(Callback): 80 | '''Learning rate scheduler. 81 | 82 | # Arguments 83 | schedule: a function that takes an epoch index as input 84 | (integer, indexed from 0) and returns a new 85 | learning rate as output (float). 86 | ''' 87 | def __init__(self, decay, every_n=1, verbose=0): 88 | Callback.__init__(self) 89 | self.decay = decay 90 | self.every_n = every_n 91 | self.verbose = verbose 92 | 93 | def on_epoch_end(self, epoch, logs={}): 94 | if not (epoch and epoch % self.every_n == 0): 95 | return 96 | 97 | assert hasattr(self.model.optimizer, 'lr'), \ 98 | 'Optimizer must have a "lr" attribute.' 99 | current_lr = K.get_value(self.model.optimizer.lr) 100 | new_lr = current_lr * self.decay 101 | if self.verbose > 0: 102 | print(' \nEpoch %05d: reducing learning rate' % (epoch)) 103 | sys.stderr.write('new lr: %.5f\n' % new_lr) 104 | K.set_value(self.model.optimizer.lr, new_lr) 105 | -------------------------------------------------------------------------------- /current_ensemble.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import sys 3 | from data import load_test_data 4 | from u_model import get_unet 5 | from keras.optimizers import Adam 6 | from train import preprocess, Learner 7 | #aug 8 | from u_model import IMG_COLS as img_cols, IMG_ROWS as img_rows 9 | from keras.preprocessing.image import flip_axis, random_channel_shift 10 | 11 | curry = lambda func, *args, **kw:\ 12 | lambda *p, **n:\ 13 | func(*args + p, **dict(kw.items() + n.items())) 14 | from keras.preprocessing.image import apply_transform, transform_matrix_offset_center 15 | 16 | 17 | def zoom(x, zoom_range, row_index=1, col_index=2, channel_index=0, 18 | fill_mode='nearest', cval=0.): 19 | zx, zy = zoom_range 20 | zoom_matrix = np.array([[zx, 0, 0], 21 | [0, zy, 0], 22 | [0, 0, 1]]) 23 | 24 | h, w = x.shape[row_index], x.shape[col_index] 25 | transform_matrix = transform_matrix_offset_center(zoom_matrix, h, w) 26 | x = apply_transform(x, transform_matrix, channel_index, fill_mode, cval) 27 | return x 28 | 29 | 30 | transforms = ( 31 | {'do': curry(flip_axis, axis=1), 'undo': curry(flip_axis, axis=1)}, 32 | {'do': curry(flip_axis, axis=2), 'undo': curry(flip_axis, axis=2)}, 33 | {'do': curry(zoom, zoom_range=(1.05, 1.05)), 'undo': curry(zoom, zoom_range=(1/1.05, 1/1.05))}, 34 | {'do': curry(zoom, zoom_range=(0.95, 0.95)), 'undo': curry(zoom, zoom_range=(1/0.95, 1/0.95))}, 35 | {'do': curry(random_channel_shift, intensity=5), 'undo': lambda x: x}, 36 | ) 37 | 38 | 39 | def run_test(): 40 | BS = 256 41 | print('Loading and preprocessing test data...') 42 | mean, std = Learner.load_meanstd() 43 | 44 | imgs_test = load_test_data() 45 | # imgs_test = imgs_test[:100] 46 | # print ('test') 47 | imgs_test = preprocess(imgs_test) 48 | 49 | imgs_test = imgs_test.astype('float32') 50 | imgs_test -= mean 51 | imgs_test /= std 52 | 53 | 54 | print ('Augment') 55 | alen, dlen = len(transforms), len(imgs_test) 56 | test_x = np.ndarray((alen, dlen, 1, img_rows, img_cols), dtype=np.float32) 57 | for i in xrange(dlen): 58 | for j, transform in enumerate(transforms): 59 | test_x[j,i] = transform['do'](imgs_test[i].copy()) 60 | # 61 | kfold = 6 62 | kfold_masks, kfold_prob = [], [] 63 | for _iter in xrange(kfold): 64 | print('Iter=%d, Loading saved weights...' % _iter) 65 | model = get_unet(Adam(0.001)) 66 | filepath = Learner.best_weight_path + '_%d.fold' % _iter 67 | print ('Loading weights from %s' % filepath) 68 | model.load_weights(filepath) 69 | # 70 | print('Predicting masks on test data...') 71 | outputs = [] 72 | asis_res = model.predict(imgs_test, batch_size=BS, verbose=1) 73 | outputs.append(asis_res) 74 | for j, transform in enumerate(transforms): 75 | t_y = model.predict(test_x[j], batch_size=BS, verbose=1) 76 | outputs.append(t_y) 77 | # 78 | print('Analyzing') 79 | test_masks = np.ndarray((dlen, 1, img_rows, img_cols), dtype=np.float32) 80 | test_probs = np.ndarray((dlen, ), dtype=np.float32) 81 | for i in xrange(dlen): 82 | masks = np.ndarray((alen+1, 1, img_rows, img_cols), dtype=np.float32) 83 | probs = np.ndarray((alen+1, ), dtype=np.float32) 84 | for j, t_y in enumerate(outputs): 85 | mask, prob = t_y[0][i], t_y[1][i] 86 | if j: 87 | mask = transforms[j-1]['undo'](mask.copy()) 88 | masks[j] = mask 89 | probs[j] = prob 90 | # 91 | test_masks[i] = np.mean(masks, 0) 92 | test_probs[i] = np.mean(probs) 93 | kfold_masks.append(test_masks) 94 | kfold_prob.append(test_probs) 95 | 96 | print 'Summing results of ensemble' 97 | # 98 | res_masks = np.ndarray((dlen, 1, img_rows, img_cols), dtype=np.float32) 99 | res_probs = np.ndarray((dlen, ), dtype=np.float32) 100 | for i in xrange(dlen): 101 | masks = np.ndarray((kfold, 1, img_rows, img_cols), dtype=np.float32) 102 | probs = np.ndarray((kfold, ), dtype=np.float32) 103 | for k in xrange(kfold): 104 | masks[k] = kfold_masks[k][i] 105 | probs[k] = kfold_prob[k][i] 106 | res_masks[i] = np.mean(masks, 0) 107 | res_probs[i] = np.mean(probs) 108 | 109 | 110 | print('Saving ') 111 | np.save(Learner.test_mask_res, res_masks) 112 | np.save(Learner.test_mask_exist_res, res_probs) 113 | 114 | 115 | def main(): 116 | run_test() 117 | 118 | if __name__ == '__main__': 119 | sys.exit(main()) 120 | 121 | -------------------------------------------------------------------------------- /train_kfold.py: -------------------------------------------------------------------------------- 1 | from optparse import OptionParser 2 | import cv2, sys, os, shutil, random 3 | import numpy as np 4 | from keras.optimizers import Adam, SGD, RMSprop 5 | from keras.callbacks import ModelCheckpoint, EarlyStopping 6 | from keras.preprocessing.image import flip_axis, random_channel_shift 7 | from keras.engine.training import slice_X 8 | from keras_plus import LearningRateDecay 9 | from u_model import get_unet, IMG_COLS as img_cols, IMG_ROWS as img_rows 10 | from data import load_train_data, load_test_data, load_patient_num 11 | from augmentation import CustomImageDataGenerator 12 | from augmentation import random_zoom, elastic_transform, random_rotation 13 | from utils import save_pickle, load_pickle, count_enum 14 | from sklearn.cross_validation import KFold 15 | 16 | _dir = os.path.join(os.path.realpath(os.path.dirname(__file__)), '') 17 | 18 | 19 | 20 | def preprocess(imgs, to_rows=None, to_cols=None): 21 | if to_rows is None or to_cols is None: 22 | to_rows = img_rows 23 | to_cols = img_cols 24 | imgs_p = np.ndarray((imgs.shape[0], imgs.shape[1], to_rows, to_cols), dtype=np.uint8) 25 | for i in xrange(imgs.shape[0]): 26 | imgs_p[i, 0] = cv2.resize(imgs[i, 0], (to_cols, to_rows), interpolation=cv2.INTER_CUBIC) 27 | return imgs_p 28 | 29 | class Learner(object): 30 | 31 | suffix = '' 32 | res_dir = os.path.join(_dir, 'res' + suffix) 33 | best_weight_path = os.path.join(res_dir, 'unet.hdf5') 34 | test_mask_res = os.path.join(res_dir, 'imgs_mask_test.npy') 35 | test_mask_exist_res = os.path.join(res_dir, 'imgs_mask_exist_test.npy') 36 | meanstd_path = os.path.join(res_dir, 'meanstd.dump') 37 | valid_data_path = os.path.join(res_dir, 'valid.npy') 38 | tensorboard_dir = os.path.join(res_dir, 'tb') 39 | 40 | def __init__(self, model_func, validation_split): 41 | self.model_func = model_func 42 | self.validation_split = validation_split 43 | self.__iter_res_dir = os.path.join(self.res_dir, 'res_iter') 44 | self.__iter_res_file = os.path.join(self.__iter_res_dir, '{epoch:02d}-{val_loss:.4f}.unet.hdf5') 45 | 46 | def _dir_init(self): 47 | if not os.path.exists(self.res_dir): 48 | os.mkdir(self.res_dir) 49 | #iter clean 50 | if os.path.exists(self.__iter_res_dir): 51 | shutil.rmtree(self.__iter_res_dir) 52 | os.mkdir(self.__iter_res_dir) 53 | 54 | def save_meanstd(self): 55 | data = [self.mean, self.std] 56 | save_pickle(self.meanstd_path, data) 57 | 58 | @classmethod 59 | def load_meanstd(cls): 60 | print ('Load meanstd from %s' % cls.meanstd_path) 61 | mean, std = load_pickle(cls.meanstd_path) 62 | return mean, std 63 | 64 | @classmethod 65 | def save_valid_idx(cls, idx): 66 | save_pickle(cls.valid_data_path, idx) 67 | 68 | @classmethod 69 | def load_valid_idx(cls): 70 | return load_pickle(cls.valid_data_path) 71 | 72 | def _init_mean_std(self, data): 73 | data = data.astype('float32') 74 | self.mean, self.std = np.mean(data), np.std(data) 75 | self.save_meanstd() 76 | return data 77 | 78 | def get_object_existance(self, mask_array): 79 | return np.array([int(np.sum(mask_array[i, 0]) > 0) for i in xrange(len(mask_array))]) 80 | 81 | def standartize(self, array, to_float=False): 82 | if to_float: 83 | array = array.astype('float32') 84 | if self.mean is None or self.std is None: 85 | raise ValueError, 'No mean/std is initialised' 86 | 87 | array -= self.mean 88 | array /= self.std 89 | return array 90 | 91 | @classmethod 92 | def norm_mask(cls, mask_array): 93 | mask_array = mask_array.astype('float32') 94 | mask_array /= 255.0 95 | return mask_array 96 | 97 | @classmethod 98 | def shuffle_train(cls, data, mask): 99 | perm = np.random.permutation(len(data)) 100 | data = data[perm] 101 | mask = mask[perm] 102 | return data, mask 103 | 104 | def __pretrain_model_load(self, model, pretrained_path): 105 | if pretrained_path is not None: 106 | if not os.path.exists(pretrained_path): 107 | raise ValueError, 'No such pre-trained path exists' 108 | model.load_weights(pretrained_path) 109 | 110 | 111 | def augmentation(self, X, Y): 112 | print('Augmentation model...') 113 | total = len(X) 114 | x_train, y_train = [], [] 115 | 116 | for i in xrange(total): 117 | if i % 100 == 0: 118 | print ('Aug', i) 119 | x, y = X[i], Y[i] 120 | #standart 121 | x_train.append(x) 122 | y_train.append(y) 123 | 124 | # for _ in xrange(1): 125 | # _x, _y = elastic_transform(x[0], y[0], 100, 20) 126 | # x_train.append(_x.reshape((1,) + _x.shape)) 127 | # y_train.append(_y.reshape((1,) + _y.shape)) 128 | 129 | #flip x 130 | x_train.append(flip_axis(x, 2)) 131 | y_train.append(flip_axis(y, 2)) 132 | #flip y 133 | x_train.append(flip_axis(x, 1)) 134 | y_train.append(flip_axis(y, 1)) 135 | #continue 136 | #zoom 137 | for _ in xrange(1): 138 | _x, _y = random_zoom(x, y, (0.9, 1.1)) 139 | x_train.append(_x) 140 | y_train.append(_y) 141 | for _ in xrange(0): 142 | _x, _y = random_rotation(x, y, 5) 143 | x_train.append(_x) 144 | y_train.append(_y) 145 | #intentsity 146 | for _ in xrange(1): 147 | _x = random_channel_shift(x, 5.0) 148 | x_train.append(_x) 149 | y_train.append(y) 150 | 151 | x_train = np.array(x_train) 152 | y_train = np.array(y_train) 153 | return x_train, y_train 154 | 155 | def fit(self, x_train, y_train, nfolds=8): 156 | print('Creating and compiling and fitting model...') 157 | print('Shape:', x_train.shape) 158 | random_state = 51 159 | kf = KFold(len(x_train), n_folds=nfolds, shuffle=True, random_state=random_state) 160 | for i, (train_index, test_index) in enumerate(kf): 161 | print 'Fold %d' % i 162 | X_train, X_valid = x_train[train_index], x_train[test_index] 163 | Y_train, Y_valid = y_train[train_index], y_train[test_index] 164 | Y_valid_2 = self.get_object_existance(Y_valid) 165 | X_train, Y_train = self.augmentation(X_train, Y_train) 166 | Y_train_2 = self.get_object_existance(Y_train) 167 | # 168 | optimizer = Adam(lr=0.0045) 169 | model = self.model_func(optimizer) 170 | model_checkpoint = ModelCheckpoint(self.__iter_res_file + '_%d.fold' % i, monitor='val_loss') 171 | model_save_best = ModelCheckpoint(self.best_weight_path + '_%d.fold' % i, monitor='val_loss', 172 | save_best_only=True) 173 | early_s = EarlyStopping(monitor='val_loss', patience=8, verbose=1) 174 | # 175 | model.fit( 176 | X_train, [Y_train, Y_train_2], 177 | validation_data=(X_valid, [Y_valid, Y_valid_2]), 178 | batch_size=128, nb_epoch=40, 179 | verbose=1, shuffle=True, 180 | callbacks=[model_save_best, model_checkpoint, early_s] 181 | ) 182 | 183 | #augment 184 | return model 185 | 186 | def train_and_predict(self, pretrained_path=None): 187 | self._dir_init() 188 | print('Loading and preprocessing and standarize train data...') 189 | imgs_train, imgs_mask_train = load_train_data() 190 | imgs_train = preprocess(imgs_train) 191 | imgs_mask_train = preprocess(imgs_mask_train) 192 | imgs_mask_train = self.norm_mask(imgs_mask_train) 193 | 194 | self._init_mean_std(imgs_train) 195 | imgs_train = self.standartize(imgs_train, True) 196 | self.fit(imgs_train, imgs_mask_train) 197 | 198 | 199 | def main(): 200 | parser = OptionParser() 201 | parser.add_option("-s", "--suffix", action='store', type='str', dest='suffix', default = None) 202 | parser.add_option("-m", "--model_name", action='store', type='str', dest='model_name', default = 'u_model') 203 | # 204 | options, _ = parser.parse_args() 205 | suffix = options.suffix 206 | model_name = options.model_name 207 | if model_name is None: 208 | raise ValueError, 'model_name is not defined' 209 | # if suffix is None: 210 | # raise ValueError, 'Please specify suffix option' 211 | # print ('Suffix: "%s"' % suffix ) 212 | # 213 | import imp 214 | model_ = imp.load_source('model_', model_name + '.py') 215 | model_func = model_.get_unet 216 | # 217 | lr = Learner(model_func, validation_split=0.2) 218 | lr.train_and_predict(pretrained_path=None) 219 | print ('Results in ', lr.res_dir) 220 | 221 | if __name__ == '__main__': 222 | sys.exit(main()) 223 | -------------------------------------------------------------------------------- /u_model.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from keras.models import Model 3 | from keras.layers import Input, merge, Convolution2D, MaxPooling2D, UpSampling2D, Dense 4 | from keras.layers import BatchNormalization, Dropout, Flatten, Lambda 5 | from keras.layers.advanced_activations import ELU, LeakyReLU 6 | from metric import dice_coef, dice_coef_loss 7 | 8 | IMG_ROWS, IMG_COLS = 80, 112 9 | 10 | def _shortcut(_input, residual): 11 | stride_width = _input._keras_shape[2] / residual._keras_shape[2] 12 | stride_height = _input._keras_shape[3] / residual._keras_shape[3] 13 | equal_channels = residual._keras_shape[1] == _input._keras_shape[1] 14 | 15 | shortcut = _input 16 | # 1 X 1 conv if shape is different. Else identity. 17 | if stride_width > 1 or stride_height > 1 or not equal_channels: 18 | shortcut = Convolution2D(nb_filter=residual._keras_shape[1], nb_row=1, nb_col=1, 19 | subsample=(stride_width, stride_height), 20 | init="he_normal", border_mode="valid")(_input) 21 | 22 | return merge([shortcut, residual], mode="sum") 23 | 24 | 25 | def inception_block(inputs, depth, batch_mode=0, splitted=False, activation='relu'): 26 | assert depth % 16 == 0 27 | actv = activation == 'relu' and (lambda: LeakyReLU(0.0)) or activation == 'elu' and (lambda: ELU(1.0)) or None 28 | 29 | c1_1 = Convolution2D(depth/4, 1, 1, init='he_normal', border_mode='same')(inputs) 30 | 31 | c2_1 = Convolution2D(depth/8*3, 1, 1, init='he_normal', border_mode='same')(inputs) 32 | c2_1 = actv()(c2_1) 33 | if splitted: 34 | c2_2 = Convolution2D(depth/2, 1, 3, init='he_normal', border_mode='same')(c2_1) 35 | c2_2 = BatchNormalization(mode=batch_mode, axis=1)(c2_2) 36 | c2_2 = actv()(c2_2) 37 | c2_3 = Convolution2D(depth/2, 3, 1, init='he_normal', border_mode='same')(c2_2) 38 | else: 39 | c2_3 = Convolution2D(depth/2, 3, 3, init='he_normal', border_mode='same')(c2_1) 40 | 41 | c3_1 = Convolution2D(depth/16, 1, 1, init='he_normal', border_mode='same')(inputs) 42 | #missed batch norm 43 | c3_1 = actv()(c3_1) 44 | if splitted: 45 | c3_2 = Convolution2D(depth/8, 1, 5, init='he_normal', border_mode='same')(c3_1) 46 | c3_2 = BatchNormalization(mode=batch_mode, axis=1)(c3_2) 47 | c3_2 = actv()(c3_2) 48 | c3_3 = Convolution2D(depth/8, 5, 1, init='he_normal', border_mode='same')(c3_2) 49 | else: 50 | c3_3 = Convolution2D(depth/8, 5, 5, init='he_normal', border_mode='same')(c3_1) 51 | 52 | p4_1 = MaxPooling2D(pool_size=(3,3), strides=(1,1), border_mode='same')(inputs) 53 | c4_2 = Convolution2D(depth/8, 1, 1, init='he_normal', border_mode='same')(p4_1) 54 | 55 | res = merge([c1_1, c2_3, c3_3, c4_2], mode='concat', concat_axis=1) 56 | res = BatchNormalization(mode=batch_mode, axis=1)(res) 57 | res = actv()(res) 58 | return res 59 | 60 | 61 | def rblock(inputs, num, depth, scale=0.1): 62 | residual = Convolution2D(depth, num, num, border_mode='same')(inputs) 63 | residual = BatchNormalization(mode=2, axis=1)(residual) 64 | residual = Lambda(lambda x: x*scale)(residual) 65 | res = _shortcut(inputs, residual) 66 | return ELU()(res) 67 | 68 | 69 | def NConvolution2D(nb_filter, nb_row, nb_col, border_mode='same', subsample=(1, 1)): 70 | def f(_input): 71 | conv = Convolution2D(nb_filter=nb_filter, nb_row=nb_row, nb_col=nb_col, subsample=subsample, 72 | border_mode=border_mode)(_input) 73 | norm = BatchNormalization(mode=2, axis=1)(conv) 74 | return ELU()(norm) 75 | 76 | return f 77 | 78 | def BNA(_input): 79 | inputs_norm = BatchNormalization(mode=2, axis=1)(_input) 80 | return ELU()(inputs_norm) 81 | 82 | def reduction_a(inputs, k=64, l=64, m=96, n=96): 83 | "35x35 -> 17x17" 84 | inputs_norm = BNA(inputs) 85 | pool1 = MaxPooling2D((3,3), strides=(2,2), border_mode='same')(inputs_norm) 86 | 87 | conv2 = Convolution2D(n, 3, 3, subsample=(2,2), border_mode='same')(inputs_norm) 88 | 89 | conv3_1 = NConvolution2D(k, 1, 1, subsample=(1,1), border_mode='same')(inputs_norm) 90 | conv3_2 = NConvolution2D(l, 3, 3, subsample=(1,1), border_mode='same')(conv3_1) 91 | conv3_2 = Convolution2D(m, 3, 3, subsample=(2,2), border_mode='same')(conv3_2) 92 | 93 | res = merge([pool1, conv2, conv3_2], mode='concat', concat_axis=1) 94 | return res 95 | 96 | 97 | def reduction_b(inputs): 98 | "17x17 -> 8x8" 99 | inputs_norm = BNA(inputs) 100 | pool1 = MaxPooling2D((3,3), strides=(2,2), border_mode='same')(inputs_norm) 101 | # 102 | conv2_1 = NConvolution2D(64, 1, 1, subsample=(1,1), border_mode='same')(inputs_norm) 103 | conv2_2 = Convolution2D(96, 3, 3, subsample=(2,2), border_mode='same')(conv2_1) 104 | # 105 | conv3_1 = NConvolution2D(64, 1, 1, subsample=(1,1), border_mode='same')(inputs_norm) 106 | conv3_2 = Convolution2D(72, 3, 3, subsample=(2,2), border_mode='same')(conv3_1) 107 | # 108 | conv4_1 = NConvolution2D(64, 1, 1, subsample=(1,1), border_mode='same')(inputs_norm) 109 | conv4_2 = NConvolution2D(72, 3, 3, subsample=(1,1), border_mode='same')(conv4_1) 110 | conv4_3 = Convolution2D(80, 3, 3, subsample=(2,2), border_mode='same')(conv4_2) 111 | # 112 | res = merge([pool1, conv2_2, conv3_2, conv4_3], mode='concat', concat_axis=1) 113 | return res 114 | 115 | 116 | 117 | 118 | def get_unet_inception_2head(optimizer): 119 | splitted = True 120 | act = 'elu' 121 | 122 | inputs = Input((1, IMG_ROWS, IMG_COLS), name='main_input') 123 | conv1 = inception_block(inputs, 32, batch_mode=2, splitted=splitted, activation=act) 124 | #conv1 = inception_block(conv1, 32, batch_mode=2, splitted=splitted, activation=act) 125 | 126 | #pool1 = MaxPooling2D(pool_size=(2, 2))(conv1) 127 | pool1 = NConvolution2D(32, 3, 3, border_mode='same', subsample=(2,2))(conv1) 128 | pool1 = Dropout(0.5)(pool1) 129 | 130 | conv2 = inception_block(pool1, 64, batch_mode=2, splitted=splitted, activation=act) 131 | #pool2 = MaxPooling2D(pool_size=(2, 2))(conv2) 132 | pool2 = NConvolution2D(64, 3, 3, border_mode='same', subsample=(2,2))(conv2) 133 | pool2 = Dropout(0.5)(pool2) 134 | 135 | conv3 = inception_block(pool2, 128, batch_mode=2, splitted=splitted, activation=act) 136 | #pool3 = MaxPooling2D(pool_size=(2, 2))(conv3) 137 | pool3 = NConvolution2D(128, 3, 3, border_mode='same', subsample=(2,2))(conv3) 138 | pool3 = Dropout(0.5)(pool3) 139 | 140 | conv4 = inception_block(pool3, 256, batch_mode=2, splitted=splitted, activation=act) 141 | #pool4 = MaxPooling2D(pool_size=(2, 2))(conv4) 142 | pool4 = NConvolution2D(256, 3, 3, border_mode='same', subsample=(2,2))(conv4) 143 | pool4 = Dropout(0.5)(pool4) 144 | 145 | conv5 = inception_block(pool4, 512, batch_mode=2, splitted=splitted, activation=act) 146 | #conv5 = inception_block(conv5, 512, batch_mode=2, splitted=splitted, activation=act) 147 | conv5 = Dropout(0.5)(conv5) 148 | 149 | # 150 | pre = Convolution2D(1, 1, 1, init='he_normal', activation='sigmoid')(conv5) 151 | pre = Flatten()(pre) 152 | aux_out = Dense(1, activation='sigmoid', name='aux_output')(pre) 153 | # 154 | 155 | after_conv4 = rblock(conv4, 1, 256) 156 | up6 = merge([UpSampling2D(size=(2, 2))(conv5), after_conv4], mode='concat', concat_axis=1) 157 | conv6 = inception_block(up6, 256, batch_mode=2, splitted=splitted, activation=act) 158 | conv6 = Dropout(0.5)(conv6) 159 | 160 | after_conv3 = rblock(conv3, 1, 128) 161 | up7 = merge([UpSampling2D(size=(2, 2))(conv6), after_conv3], mode='concat', concat_axis=1) 162 | conv7 = inception_block(up7, 128, batch_mode=2, splitted=splitted, activation=act) 163 | conv7 = Dropout(0.5)(conv7) 164 | 165 | after_conv2 = rblock(conv2, 1, 64) 166 | up8 = merge([UpSampling2D(size=(2, 2))(conv7), after_conv2], mode='concat', concat_axis=1) 167 | conv8 = inception_block(up8, 64, batch_mode=2, splitted=splitted, activation=act) 168 | conv8 = Dropout(0.5)(conv8) 169 | 170 | after_conv1 = rblock(conv1, 1, 32) 171 | up9 = merge([UpSampling2D(size=(2, 2))(conv8), after_conv1], mode='concat', concat_axis=1) 172 | conv9 = inception_block(up9, 32, batch_mode=2, splitted=splitted, activation=act) 173 | #conv9 = inception_block(conv9, 32, batch_mode=2, splitted=splitted, activation=act) 174 | conv9 = Dropout(0.5)(conv9) 175 | 176 | conv10 = Convolution2D(1, 1, 1, init='he_normal', activation='sigmoid', name='main_output')(conv9) 177 | #print conv10._keras_shape 178 | 179 | model = Model(input=inputs, output=[conv10, aux_out]) 180 | model.compile(optimizer=optimizer, 181 | loss={'main_output': dice_coef_loss, 'aux_output': 'binary_crossentropy'}, 182 | metrics={'main_output': dice_coef, 'aux_output': 'acc'}, 183 | loss_weights={'main_output': 1., 'aux_output': 0.5}) 184 | 185 | return model 186 | 187 | 188 | get_unet = get_unet_inception_2head 189 | 190 | def main(): 191 | from keras.optimizers import Adam, RMSprop, SGD 192 | from metric import dice_coef, dice_coef_loss 193 | import numpy as np 194 | img_rows = IMG_ROWS 195 | img_cols = IMG_COLS 196 | 197 | optimizer = RMSprop(lr=0.045, rho=0.9, epsilon=1.0) 198 | model = get_unet(Adam(lr=1e-5)) 199 | model.compile(optimizer=optimizer, loss=dice_coef_loss, metrics=[dice_coef]) 200 | 201 | x = np.random.random((1, 1,img_rows,img_cols)) 202 | res = model.predict(x, 1) 203 | print res 204 | #print 'res', res[0].shape 205 | print 'params', model.count_params() 206 | print 'layer num', len(model.layers) 207 | # 208 | 209 | 210 | if __name__ == '__main__': 211 | sys.exit(main()) 212 | 213 | -------------------------------------------------------------------------------- /augmentation.py: -------------------------------------------------------------------------------- 1 | import sys, os 2 | import numpy as np 3 | from keras.preprocessing.image import (transform_matrix_offset_center, apply_transform, Iterator, 4 | random_channel_shift, flip_axis) 5 | from scipy.ndimage.interpolation import map_coordinates 6 | from scipy.ndimage.filters import gaussian_filter 7 | 8 | 9 | _dir = os.path.join(os.path.realpath(os.path.dirname(__file__)), '') 10 | data_path = os.path.join(_dir, '../') 11 | aug_data_path = os.path.join(_dir, 'aug_data') 12 | aug_pattern = os.path.join(aug_data_path, 'train_img_%d.npy') 13 | aug_mask_pattern = os.path.join(aug_data_path, 'train_mask_%d.npy') 14 | 15 | 16 | def random_zoom(x, y, zoom_range, row_index=1, col_index=2, channel_index=0, 17 | fill_mode='nearest', cval=0.): 18 | if len(zoom_range) != 2: 19 | raise Exception('zoom_range should be a tuple or list of two floats. ' 20 | 'Received arg: ', zoom_range) 21 | 22 | if zoom_range[0] == 1 and zoom_range[1] == 1: 23 | zx, zy = 1, 1 24 | else: 25 | zx, zy = np.random.uniform(zoom_range[0], zoom_range[1], 2) 26 | zoom_matrix = np.array([[zx, 0, 0], 27 | [0, zy, 0], 28 | [0, 0, 1]]) 29 | 30 | h, w = x.shape[row_index], x.shape[col_index] 31 | transform_matrix = transform_matrix_offset_center(zoom_matrix, h, w) 32 | x = apply_transform(x, transform_matrix, channel_index, fill_mode, cval) 33 | y = apply_transform(y, transform_matrix, channel_index, fill_mode, cval) 34 | return x, y 35 | 36 | 37 | def random_rotation(x, y, rg, row_index=1, col_index=2, channel_index=0, 38 | fill_mode='nearest', cval=0.): 39 | theta = np.pi / 180 * np.random.uniform(-rg, rg) 40 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], 41 | [np.sin(theta), np.cos(theta), 0], 42 | [0, 0, 1]]) 43 | 44 | h, w = x.shape[row_index], x.shape[col_index] 45 | transform_matrix = transform_matrix_offset_center(rotation_matrix, h, w) 46 | x = apply_transform(x, transform_matrix, channel_index, fill_mode, cval) 47 | y = apply_transform(y, transform_matrix, channel_index, fill_mode, cval) 48 | return x, y 49 | 50 | 51 | def random_shear(x, y, intensity, row_index=1, col_index=2, channel_index=0, 52 | fill_mode='constant', cval=0.): 53 | shear = np.random.uniform(-intensity, intensity) 54 | shear_matrix = np.array([[1, -np.sin(shear), 0], 55 | [0, np.cos(shear), 0], 56 | [0, 0, 1]]) 57 | 58 | h, w = x.shape[row_index], x.shape[col_index] 59 | transform_matrix = transform_matrix_offset_center(shear_matrix, h, w) 60 | x = apply_transform(x, transform_matrix, channel_index, fill_mode, cval) 61 | y = apply_transform(y, transform_matrix, channel_index, fill_mode, cval) 62 | return x, y 63 | 64 | 65 | class CustomNumpyArrayIterator(Iterator): 66 | 67 | def __init__(self, X, y, image_data_generator, 68 | batch_size=32, shuffle=False, seed=None, 69 | dim_ordering='th'): 70 | self.X = X 71 | self.y = y 72 | self.image_data_generator = image_data_generator 73 | self.dim_ordering = dim_ordering 74 | super(CustomNumpyArrayIterator, self).__init__(X.shape[0], batch_size, shuffle, seed) 75 | 76 | 77 | def next(self): 78 | with self.lock: 79 | index_array, _, current_batch_size = next(self.index_generator) 80 | batch_x = np.zeros(tuple([current_batch_size] + list(self.X.shape)[1:])) 81 | batch_y_1, batch_y_2 = [], [] 82 | for i, j in enumerate(index_array): 83 | x = self.X[j] 84 | y1 = self.y[0][j] 85 | y2 = self.y[1][j] 86 | _x, _y1 = self.image_data_generator.random_transform(x.astype('float32'), y1.astype('float32')) 87 | batch_x[i] = _x 88 | batch_y_1.append(_y1) 89 | batch_y_2.append(y2) 90 | return batch_x, [np.array(batch_y_1), np.array(batch_y_2)] 91 | 92 | 93 | class CustomImageDataGenerator(object): 94 | def __init__(self, zoom_range=(1,1), channel_shift_range=0, horizontal_flip=False, vertical_flip=False, 95 | rotation_range=0, 96 | width_shift_range=0., 97 | height_shift_range=0., 98 | shear_range=0., 99 | elastic=None, 100 | ): 101 | self.zoom_range = zoom_range 102 | self.channel_shift_range = channel_shift_range 103 | self.horizontal_flip = horizontal_flip 104 | self.vertical_flip = vertical_flip 105 | self.rotation_range = rotation_range 106 | self.width_shift_range = width_shift_range 107 | self.height_shift_range = height_shift_range 108 | self.shear_range = shear_range 109 | self.elastic = elastic 110 | 111 | def random_transform(self, x, y, row_index=1, col_index=2, channel_index=0): 112 | 113 | if self.horizontal_flip: 114 | if True or np.random.random() < 0.5: 115 | x = flip_axis(x, 2) 116 | y = flip_axis(y, 2) 117 | 118 | # use composition of homographies to generate final transform that needs to be applied 119 | if self.rotation_range: 120 | theta = np.pi / 180 * np.random.uniform(-self.rotation_range, self.rotation_range) 121 | else: 122 | theta = 0 123 | rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0], 124 | [np.sin(theta), np.cos(theta), 0], 125 | [0, 0, 1]]) 126 | if self.height_shift_range: 127 | tx = np.random.uniform(-self.height_shift_range, self.height_shift_range) * x.shape[row_index] 128 | else: 129 | tx = 0 130 | 131 | if self.width_shift_range: 132 | ty = np.random.uniform(-self.width_shift_range, self.width_shift_range) * x.shape[col_index] 133 | else: 134 | ty = 0 135 | 136 | translation_matrix = np.array([[1, 0, tx], 137 | [0, 1, ty], 138 | [0, 0, 1]]) 139 | if self.shear_range: 140 | shear = np.random.uniform(-self.shear_range, self.shear_range) 141 | else: 142 | shear = 0 143 | shear_matrix = np.array([[1, -np.sin(shear), 0], 144 | [0, np.cos(shear), 0], 145 | [0, 0, 1]]) 146 | 147 | if self.zoom_range[0] == 1 and self.zoom_range[1] == 1: 148 | zx, zy = 1, 1 149 | else: 150 | zx, zy = np.random.uniform(self.zoom_range[0], self.zoom_range[1], 2) 151 | zoom_matrix = np.array([[zx, 0, 0], 152 | [0, zy, 0], 153 | [0, 0, 1]]) 154 | 155 | transform_matrix = np.dot(np.dot(np.dot(rotation_matrix, translation_matrix), shear_matrix), zoom_matrix) 156 | 157 | h, w = x.shape[row_index], x.shape[col_index] 158 | transform_matrix = transform_matrix_offset_center(transform_matrix, h, w) 159 | 160 | x = apply_transform(x, transform_matrix, channel_index, 161 | fill_mode='constant') 162 | y = apply_transform(y, transform_matrix, channel_index, 163 | fill_mode='constant') 164 | 165 | 166 | # 167 | 168 | if self.vertical_flip: 169 | if np.random.random() < 0.5: 170 | x = flip_axis(x, 1) 171 | y = flip_axis(y, 1) 172 | 173 | if self.channel_shift_range != 0: 174 | x = random_channel_shift(x, self.channel_shift_range) 175 | 176 | 177 | if self.elastic is not None: 178 | x, y = elastic_transform(x.reshape(h,w), y.reshape(h,w), *self.elastic) 179 | x, y = x.reshape(1, h, w), y.reshape(1, h, w) 180 | 181 | return x, y 182 | 183 | def flow(self, X, Y, batch_size, shuffle=True, seed=None): 184 | return CustomNumpyArrayIterator( 185 | X, Y, self, 186 | batch_size=batch_size, shuffle=shuffle, seed=seed) 187 | 188 | 189 | def elastic_transform(image, mask, alpha, sigma, alpha_affine=None, random_state=None): 190 | """Elastic deformation of images as described in [Simard2003]_ (with modifications). 191 | .. [Simard2003] Simard, Steinkraus and Platt, "Best Practices for 192 | Convolutional Neural Networks applied to Visual Document Analysis", in 193 | Proc. of the International Conference on Document Analysis and 194 | Recognition, 2003. 195 | 196 | Based on https://gist.github.com/erniejunior/601cdf56d2b424757de5 197 | """ 198 | if random_state is None: 199 | random_state = np.random.RandomState(None) 200 | 201 | shape = image.shape 202 | 203 | dx = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha 204 | dy = gaussian_filter((random_state.rand(*shape) * 2 - 1), sigma) * alpha 205 | 206 | x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0])) 207 | indices = np.reshape(y+dy, (-1, 1)), np.reshape(x+dx, (-1, 1)) 208 | 209 | 210 | res_x = map_coordinates(image, indices, order=1, mode='reflect').reshape(shape) 211 | res_y = map_coordinates(mask, indices, order=1, mode='reflect').reshape(shape) 212 | return res_x, res_y 213 | 214 | 215 | def test(): 216 | X = np.random.randint(0,100, (1000, 1, 100, 200)) 217 | YY = [np.random.randint(0,100, (1000, 1, 100, 200)), np.random.random((1000, 1))] 218 | cid = CustomImageDataGenerator(horizontal_flip=True, elastic=(100,20)) 219 | gen = cid.flow(X, YY, batch_size=64, shuffle=False) 220 | n = gen.next()[0] 221 | 222 | 223 | if __name__ == '__main__': 224 | sys.exit(test()) 225 | -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from optparse import OptionParser 3 | import cv2, sys, os, shutil, random 4 | import numpy as np 5 | from keras.optimizers import Adam 6 | from keras.callbacks import ModelCheckpoint, EarlyStopping 7 | from keras.preprocessing.image import flip_axis, random_channel_shift 8 | from keras.engine.training import slice_X 9 | from keras_plus import LearningRateDecay 10 | from u_model import get_unet, IMG_COLS as img_cols, IMG_ROWS as img_rows 11 | from data import load_train_data, load_test_data, load_patient_num 12 | from augmentation import random_zoom, elastic_transform, random_rotation 13 | from utils import save_pickle, load_pickle, count_enum 14 | 15 | _dir = os.path.join(os.path.realpath(os.path.dirname(__file__)), '') 16 | 17 | 18 | def preprocess(imgs, to_rows=None, to_cols=None): 19 | if to_rows is None or to_cols is None: 20 | to_rows = img_rows 21 | to_cols = img_cols 22 | imgs_p = np.ndarray((imgs.shape[0], imgs.shape[1], to_rows, to_cols), dtype=np.uint8) 23 | for i in xrange(imgs.shape[0]): 24 | imgs_p[i, 0] = cv2.resize(imgs[i, 0], (to_cols, to_rows), interpolation=cv2.INTER_CUBIC) 25 | return imgs_p 26 | 27 | 28 | class Learner(object): 29 | 30 | suffix = '' 31 | res_dir = os.path.join(_dir, 'res' + suffix) 32 | best_weight_path = os.path.join(res_dir, 'unet.hdf5') 33 | test_mask_res = os.path.join(res_dir, 'imgs_mask_test.npy') 34 | test_mask_exist_res = os.path.join(res_dir, 'imgs_mask_exist_test.npy') 35 | meanstd_path = os.path.join(res_dir, 'meanstd.dump') 36 | valid_data_path = os.path.join(res_dir, 'valid.npy') 37 | tensorboard_dir = os.path.join(res_dir, 'tb') 38 | 39 | def __init__(self, model_func, validation_split): 40 | self.model_func = model_func 41 | self.validation_split = validation_split 42 | self.__iter_res_dir = os.path.join(self.res_dir, 'res_iter') 43 | self.__iter_res_file = os.path.join(self.__iter_res_dir, '{epoch:02d}-{val_loss:.4f}.unet.hdf5') 44 | 45 | def _dir_init(self): 46 | if not os.path.exists(self.res_dir): 47 | os.mkdir(self.res_dir) 48 | #iter clean 49 | if os.path.exists(self.__iter_res_dir): 50 | shutil.rmtree(self.__iter_res_dir) 51 | os.mkdir(self.__iter_res_dir) 52 | 53 | def save_meanstd(self): 54 | data = [self.mean, self.std] 55 | save_pickle(self.meanstd_path, data) 56 | 57 | @classmethod 58 | def load_meanstd(cls): 59 | print ('Load meanstd from %s' % cls.meanstd_path) 60 | mean, std = load_pickle(cls.meanstd_path) 61 | return mean, std 62 | 63 | @classmethod 64 | def save_valid_idx(cls, idx): 65 | save_pickle(cls.valid_data_path, idx) 66 | 67 | @classmethod 68 | def load_valid_idx(cls): 69 | return load_pickle(cls.valid_data_path) 70 | 71 | def _init_mean_std(self, data): 72 | data = data.astype('float32') 73 | self.mean, self.std = np.mean(data), np.std(data) 74 | self.save_meanstd() 75 | return data 76 | 77 | def get_object_existance(self, mask_array): 78 | return np.array([int(np.sum(mask_array[i, 0]) > 0) for i in xrange(len(mask_array))]) 79 | 80 | def standartize(self, array, to_float=False): 81 | if to_float: 82 | array = array.astype('float32') 83 | if self.mean is None or self.std is None: 84 | raise ValueError, 'No mean/std is initialised' 85 | 86 | array -= self.mean 87 | array /= self.std 88 | return array 89 | 90 | @classmethod 91 | def norm_mask(cls, mask_array): 92 | mask_array = mask_array.astype('float32') 93 | mask_array /= 255.0 94 | return mask_array 95 | 96 | @classmethod 97 | def shuffle_train(cls, data, mask): 98 | perm = np.random.permutation(len(data)) 99 | data = data[perm] 100 | mask = mask[perm] 101 | return data, mask 102 | 103 | @classmethod 104 | def split_train_and_valid_by_patient(cls, data, mask, validation_split, shuffle=False): 105 | print('Shuffle & split...') 106 | patient_nums = load_patient_num() 107 | patient_dict = count_enum(patient_nums) 108 | pnum = len(patient_dict) 109 | val_num = int(pnum * validation_split) 110 | patients = patient_dict.keys() 111 | if shuffle: 112 | random.shuffle(patients) 113 | val_p, train_p = patients[:val_num], patients[val_num:] 114 | train_indexes = [i for i, c in enumerate(patient_nums) if c in set(train_p)] 115 | val_indexes = [i for i, c in enumerate(patient_nums) if c in set(val_p)] 116 | x_train, y_train = data[train_indexes], mask[train_indexes] 117 | x_valid, y_valid = data[val_indexes], mask[val_indexes] 118 | cls.save_valid_idx(val_indexes) 119 | print ('val patients:', len(x_valid), val_p) 120 | print ('train patients:', len(x_train), train_p) 121 | return (x_train, y_train), (x_valid, y_valid) 122 | 123 | @classmethod 124 | def split_train_and_valid(cls, data, mask, validation_split, shuffle=False): 125 | print('Shuffle & split...') 126 | if shuffle: 127 | data, mask = cls.shuffle_train(data, mask) 128 | split_at = int(len(data) * (1. - validation_split)) 129 | x_train, x_valid = (slice_X(data, 0, split_at), slice_X(data, split_at)) 130 | y_train, y_valid = (slice_X(mask, 0, split_at), slice_X(mask, split_at)) 131 | cls.save_valid_idx(range(len(data))[split_at:]) 132 | return (x_train, y_train), (x_valid, y_valid) 133 | 134 | def test(self, model, batch_size=256): 135 | print('Loading and pre-processing test data...') 136 | imgs_test = load_test_data() 137 | imgs_test = preprocess(imgs_test) 138 | imgs_test = self.standartize(imgs_test, to_float=True) 139 | 140 | print('Loading best saved weights...') 141 | model.load_weights(self.best_weight_path) 142 | print('Predicting masks on test data and saving...') 143 | imgs_mask_test = model.predict(imgs_test, batch_size=batch_size, verbose=1) 144 | 145 | np.save(self.test_mask_res, imgs_mask_test[0]) 146 | np.save(self.test_mask_exist_res, imgs_mask_test[1]) 147 | 148 | def __pretrain_model_load(self, model, pretrained_path): 149 | if pretrained_path is not None: 150 | if not os.path.exists(pretrained_path): 151 | raise ValueError, 'No such pre-trained path exists' 152 | model.load_weights(pretrained_path) 153 | 154 | 155 | def augmentation(self, X, Y): 156 | print('Augmentation model...') 157 | total = len(X) 158 | x_train, y_train = [], [] 159 | 160 | for i in xrange(total): 161 | x, y = X[i], Y[i] 162 | #standart 163 | x_train.append(x) 164 | y_train.append(y) 165 | 166 | # for _ in xrange(1): 167 | # _x, _y = elastic_transform(x[0], y[0], 100, 20) 168 | # x_train.append(_x.reshape((1,) + _x.shape)) 169 | # y_train.append(_y.reshape((1,) + _y.shape)) 170 | 171 | #flip x 172 | x_train.append(flip_axis(x, 2)) 173 | y_train.append(flip_axis(y, 2)) 174 | #flip y 175 | x_train.append(flip_axis(x, 1)) 176 | y_train.append(flip_axis(y, 1)) 177 | #continue 178 | #zoom 179 | for _ in xrange(1): 180 | _x, _y = random_zoom(x, y, (0.9, 1.1)) 181 | x_train.append(_x) 182 | y_train.append(_y) 183 | for _ in xrange(0): 184 | _x, _y = random_rotation(x, y, 5) 185 | x_train.append(_x) 186 | y_train.append(_y) 187 | #intentsity 188 | for _ in xrange(1): 189 | _x = random_channel_shift(x, 5.0) 190 | x_train.append(_x) 191 | y_train.append(y) 192 | 193 | x_train = np.array(x_train) 194 | y_train = np.array(y_train) 195 | return x_train, y_train 196 | 197 | def fit(self, x_train, y_train, x_valid, y_valid, pretrained_path): 198 | print('Creating and compiling and fitting model...') 199 | print('Shape:', x_train.shape) 200 | #second output 201 | y_train_2 = self.get_object_existance(y_train) 202 | y_valid_2 = self.get_object_existance(y_valid) 203 | 204 | #load model 205 | optimizer = Adam(lr=0.0045) 206 | model = self.model_func(optimizer) 207 | 208 | #checkpoints 209 | model_checkpoint = ModelCheckpoint(self.__iter_res_file, monitor='val_loss') 210 | model_save_best = ModelCheckpoint(self.best_weight_path, monitor='val_loss', save_best_only=True) 211 | early_s = EarlyStopping(monitor='val_loss', patience=5, verbose=1) 212 | learning_rate_adapt = LearningRateDecay(0.9, every_n=2, verbose=1) 213 | self.__pretrain_model_load(model, pretrained_path) 214 | model.fit( 215 | x_train, [y_train, y_train_2], 216 | validation_data=(x_valid, [y_valid, y_valid_2]), 217 | batch_size=128, nb_epoch=50, 218 | verbose=1, shuffle=True, 219 | callbacks=[model_save_best, model_checkpoint, early_s] 220 | ) 221 | 222 | #augment 223 | return model 224 | 225 | def train_and_predict(self, pretrained_path=None, split_random=True): 226 | self._dir_init() 227 | print('Loading and preprocessing and standarize train data...') 228 | imgs_train, imgs_mask_train = load_train_data() 229 | 230 | imgs_train = preprocess(imgs_train) 231 | 232 | imgs_mask_train = preprocess(imgs_mask_train) 233 | 234 | imgs_mask_train = self.norm_mask(imgs_mask_train) 235 | 236 | split_func = split_random and self.split_train_and_valid or self.split_train_and_valid_by_patient 237 | (x_train, y_train), (x_valid, y_valid) = split_func(imgs_train, imgs_mask_train, 238 | validation_split=self.validation_split) 239 | self._init_mean_std(x_train) 240 | x_train = self.standartize(x_train, True) 241 | x_valid = self.standartize(x_valid, True) 242 | #augmentation 243 | x_train, y_train = self.augmentation(x_train, y_train) 244 | #fit 245 | model = self.fit(x_train, y_train, x_valid, y_valid, pretrained_path) 246 | #test 247 | self.test(model) 248 | 249 | 250 | def main(): 251 | parser = OptionParser() 252 | parser.add_option("-s", "--split_random", action='store', type='int', dest='split_random', default = 1) 253 | parser.add_option("-m", "--model_name", action='store', type='str', dest='model_name', default = 'u_model') 254 | # 255 | options, _ = parser.parse_args() 256 | split_random = options.split_random 257 | model_name = options.model_name 258 | if model_name is None: 259 | raise ValueError, 'model_name is not defined' 260 | # 261 | import imp 262 | model_ = imp.load_source('model_', model_name + '.py') 263 | model_func = model_.get_unet 264 | # 265 | lr = Learner(model_func, validation_split=0.2) 266 | lr.train_and_predict(pretrained_path=None, split_random=split_random) 267 | print ('Results in ', lr.res_dir) 268 | 269 | if __name__ == '__main__': 270 | sys.exit(main()) 271 | -------------------------------------------------------------------------------- /train_generator.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | from optparse import OptionParser 3 | import cv2, sys, os, shutil, random 4 | import numpy as np 5 | from keras.optimizers import Adam, SGD, RMSprop 6 | from keras.callbacks import ModelCheckpoint, EarlyStopping 7 | from keras.preprocessing.image import flip_axis, random_channel_shift 8 | from keras.engine.training import slice_X 9 | from keras_plus import LearningRateDecay 10 | from u_model import get_unet, IMG_COLS as img_cols, IMG_ROWS as img_rows 11 | from data import load_train_data, load_test_data, load_patient_num 12 | from augmentation import CustomImageDataGenerator 13 | from augmentation import random_zoom, elastic_transform, load_aug 14 | from utils import save_pickle, load_pickle, count_enum 15 | 16 | _dir = os.path.join(os.path.realpath(os.path.dirname(__file__)), '') 17 | 18 | 19 | 20 | def preprocess(imgs, to_rows=None, to_cols=None): 21 | if to_rows is None or to_cols is None: 22 | to_rows = img_rows 23 | to_cols = img_cols 24 | imgs_p = np.ndarray((imgs.shape[0], imgs.shape[1], to_rows, to_cols), dtype=np.uint8) 25 | for i in xrange(imgs.shape[0]): 26 | imgs_p[i, 0] = cv2.resize(imgs[i, 0], (to_cols, to_rows), interpolation=cv2.INTER_CUBIC) 27 | return imgs_p 28 | 29 | class Learner(object): 30 | 31 | suffix = '' 32 | res_dir = os.path.join(_dir, 'res' + suffix) 33 | best_weight_path = os.path.join(res_dir, 'unet.hdf5') 34 | test_mask_res = os.path.join(res_dir, 'imgs_mask_test.npy') 35 | test_mask_exist_res = os.path.join(res_dir, 'imgs_mask_exist_test.npy') 36 | meanstd_path = os.path.join(res_dir, 'meanstd.dump') 37 | valid_data_path = os.path.join(res_dir, 'valid.npy') 38 | tensorboard_dir = os.path.join(res_dir, 'tb') 39 | 40 | def __init__(self, model_func, validation_split): 41 | self.model_func = model_func 42 | self.validation_split = validation_split 43 | self.__iter_res_dir = os.path.join(self.res_dir, 'res_iter') 44 | self.__iter_res_file = os.path.join(self.__iter_res_dir, '{epoch:02d}-{val_loss:.4f}.unet.hdf5') 45 | 46 | def _dir_init(self): 47 | if not os.path.exists(self.res_dir): 48 | os.mkdir(self.res_dir) 49 | #iter clean 50 | if os.path.exists(self.__iter_res_dir): 51 | shutil.rmtree(self.__iter_res_dir) 52 | os.mkdir(self.__iter_res_dir) 53 | 54 | def save_meanstd(self): 55 | data = [self.mean, self.std] 56 | save_pickle(self.meanstd_path, data) 57 | 58 | @classmethod 59 | def load_meanstd(cls): 60 | print ('Load meanstd from %s' % cls.meanstd_path) 61 | mean, std = load_pickle(cls.meanstd_path) 62 | return mean, std 63 | 64 | @classmethod 65 | def save_valid_idx(cls, idx): 66 | save_pickle(cls.valid_data_path, idx) 67 | 68 | @classmethod 69 | def load_valid_idx(cls): 70 | return load_pickle(cls.valid_data_path) 71 | 72 | def _init_mean_std(self, data): 73 | data = data.astype('float32') 74 | self.mean, self.std = np.mean(data), np.std(data) 75 | self.save_meanstd() 76 | return data 77 | 78 | def get_object_existance(self, mask_array): 79 | return np.array([int(np.sum(mask_array[i, 0]) > 0) for i in xrange(len(mask_array))]) 80 | 81 | def standartize(self, array, to_float=False): 82 | if to_float: 83 | array = array.astype('float32') 84 | if self.mean is None or self.std is None: 85 | raise ValueError, 'No mean/std is initialised' 86 | 87 | array -= self.mean 88 | array /= self.std 89 | return array 90 | 91 | @classmethod 92 | def norm_mask(cls, mask_array): 93 | mask_array = mask_array.astype('float32') 94 | mask_array /= 255.0 95 | return mask_array 96 | 97 | @classmethod 98 | def shuffle_train(cls, data, mask): 99 | perm = np.random.permutation(len(data)) 100 | data = data[perm] 101 | mask = mask[perm] 102 | return data, mask 103 | 104 | @classmethod 105 | def split_train_and_valid_by_patient(cls, data, mask, validation_split, shuffle=False): 106 | print('Shuffle & split...') 107 | patient_nums = load_patient_num() 108 | patient_dict = count_enum(patient_nums) 109 | pnum = len(patient_dict) 110 | val_num = int(pnum * validation_split) 111 | patients = patient_dict.keys() 112 | if shuffle: 113 | random.shuffle(patients) 114 | val_p, train_p = patients[:val_num], patients[val_num:] 115 | train_indexes = [i for i, c in enumerate(patient_nums) if c in set(train_p)] 116 | val_indexes = [i for i, c in enumerate(patient_nums) if c in set(val_p)] 117 | x_train, y_train = data[train_indexes], mask[train_indexes] 118 | x_valid, y_valid = data[val_indexes], mask[val_indexes] 119 | cls.save_valid_idx(val_indexes) 120 | print ('val patients:', len(x_valid), val_p) 121 | print ('train patients:', len(x_train), train_p) 122 | return (x_train, y_train), (x_valid, y_valid) 123 | 124 | @classmethod 125 | def split_train_and_valid(cls, data, mask, validation_split, shuffle=False): 126 | print('Shuffle & split...') 127 | if shuffle: 128 | data, mask = cls.shuffle_train(data, mask) 129 | split_at = int(len(data) * (1. - validation_split)) 130 | x_train, x_valid = (slice_X(data, 0, split_at), slice_X(data, split_at)) 131 | y_train, y_valid = (slice_X(mask, 0, split_at), slice_X(mask, split_at)) 132 | cls.save_valid_idx(range(len(data))[split_at:]) 133 | return (x_train, y_train), (x_valid, y_valid) 134 | 135 | def test(self, model, batch_size=256): 136 | print('Loading and pre-processing test data...') 137 | imgs_test = load_test_data() 138 | imgs_test = preprocess(imgs_test) 139 | imgs_test = self.standartize(imgs_test, to_float=True) 140 | 141 | print('Loading best saved weights...') 142 | model.load_weights(self.best_weight_path) 143 | print('Predicting masks on test data and saving...') 144 | imgs_mask_test = model.predict(imgs_test, batch_size=batch_size, verbose=1) 145 | 146 | np.save(self.test_mask_res, imgs_mask_test[0]) 147 | np.save(self.test_mask_exist_res, imgs_mask_test[1]) 148 | 149 | def __pretrain_model_load(self, model, pretrained_path): 150 | if pretrained_path is not None: 151 | if not os.path.exists(pretrained_path): 152 | raise ValueError, 'No such pre-trained path exists' 153 | model.load_weights(pretrained_path) 154 | 155 | 156 | def augmentation(self, X, Y): 157 | print('Augmentation model...') 158 | total = len(X) 159 | x_train, y_train = [], [] 160 | 161 | for i in xrange(total): 162 | if i % 100 == 0: 163 | print ('Aug', i) 164 | x, y = X[i], Y[i] 165 | #standart 166 | x_train.append(x) 167 | y_train.append(y) 168 | 169 | for _ in xrange(2): 170 | _x, _y = elastic_transform(x[0], y[0], 100, 20) 171 | x_train.append(_x.reshape((1,) + _x.shape)) 172 | y_train.append(_y.reshape((1,) + _y.shape)) 173 | 174 | #flip x 175 | x_train.append(flip_axis(x, 2)) 176 | y_train.append(flip_axis(y, 2)) 177 | #flip y 178 | x_train.append(flip_axis(x, 1)) 179 | y_train.append(flip_axis(y, 1)) 180 | continue 181 | #zoom 182 | for _ in xrange(1): 183 | _x, _y = random_zoom(x, y, (0.9, 1.1)) 184 | x_train.append(_x) 185 | y_train.append(_y) 186 | #intentsity 187 | for _ in xrange(1): 188 | _x = random_channel_shift(x, 5.0) 189 | x_train.append(_x) 190 | y_train.append(y) 191 | 192 | # for j in xrange(5): 193 | # xs, ys = load_aug(j) 194 | # ys = self.norm_mask(ys) 195 | # (xn, yn), _ = self.split_train_and_valid_by_patient(xs, ys, validation_split=self.validation_split, shuffle=False) 196 | # for i in xrange(len(xn)): 197 | # x_train.append(xn[i]) 198 | # y_train.append(yn[i]) 199 | 200 | x_train = np.array(x_train) 201 | y_train = np.array(y_train) 202 | return x_train, y_train 203 | 204 | def fit(self, x_train, y_train, x_valid, y_valid, pretrained_path): 205 | print('Creating and compiling and fitting model...') 206 | print('Shape:', x_train.shape) 207 | #second output 208 | y_train_2 = self.get_object_existance(y_train) 209 | y_valid_2 = self.get_object_existance(y_valid) 210 | 211 | #load model 212 | optimizer = Adam(lr=0.0045) 213 | #model = get_unet(optimizer) 214 | model = self.model_func(optimizer) 215 | 216 | #checkpoints 217 | model_checkpoint = ModelCheckpoint(self.__iter_res_file, monitor='val_loss') 218 | model_save_best = ModelCheckpoint(self.best_weight_path, monitor='val_loss', save_best_only=True) 219 | early_s = EarlyStopping(monitor='val_loss', patience=10, verbose=1) 220 | #tb = TensorBoard(self.tensorboard_dir, histogram_freq=2, write_graph=True) 221 | #learning_rate_adapt = AdvancedLearnignRateScheduler(monitor='val_loss', patience=1, verbose=1, mode='min', decayRatio=0.5) 222 | learning_rate_adapt = LearningRateDecay(0.95, every_n=4, verbose=1) 223 | self.__pretrain_model_load(model, pretrained_path) 224 | #augment 225 | datagen = CustomImageDataGenerator(zoom_range=(0.9,1.1), 226 | horizontal_flip=True, 227 | vertical_flip=False, 228 | # rotation_range=5, 229 | channel_shift_range=5.0, 230 | elastic=None #(100, 20) 231 | ) 232 | # #fit 233 | model.fit_generator(datagen.flow(x_train, [y_train, y_train_2], batch_size=64), 234 | samples_per_epoch=len(x_train), 235 | nb_epoch=250, 236 | verbose=1, 237 | callbacks=[model_save_best, model_checkpoint, early_s, learning_rate_adapt], 238 | validation_data=(x_valid, [y_valid, y_valid_2]) 239 | ) 240 | return model 241 | 242 | def train_and_predict(self, pretrained_path=None, split_random=True): 243 | self._dir_init() 244 | print('Loading and preprocessing and standarize train data...') 245 | imgs_train, imgs_mask_train = load_train_data() 246 | 247 | imgs_train = preprocess(imgs_train) 248 | 249 | imgs_mask_train = preprocess(imgs_mask_train) 250 | 251 | imgs_mask_train = self.norm_mask(imgs_mask_train) 252 | #imgs_train = self.norm_mask(imgs_train) /255 253 | #shuffle and split 254 | split_func = split_random and self.split_train_and_valid or self.split_train_and_valid_by_patient 255 | (x_train, y_train), (x_valid, y_valid) = split_func(imgs_train, imgs_mask_train, 256 | validation_split=self.validation_split) 257 | self._init_mean_std(x_train) 258 | x_train = self.standartize(x_train, True) 259 | x_valid = self.standartize(x_valid, True) 260 | #fit 261 | model = self.fit(x_train, y_train, x_valid, y_valid, pretrained_path) 262 | #test 263 | self.test(model) 264 | 265 | 266 | def main(): 267 | parser = OptionParser() 268 | parser.add_option("-m", "--model_name", action='store', type='str', dest='model_name', default = 'u_model') 269 | parser.add_option("-s", "--split_random", action='store', type='int', dest='split_random', default = 1) 270 | # 271 | options, _ = parser.parse_args() 272 | model_name = options.model_name 273 | split_random = options.split_random 274 | if model_name is None: 275 | raise ValueError, 'model_name is not defined' 276 | # 277 | import imp 278 | model_ = imp.load_source('model_', model_name + '.py') 279 | model_func = model_.get_unet 280 | # 281 | lr = Learner(model_func, validation_split=0.2) 282 | lr.train_and_predict(pretrained_path=None, split_random=split_random) 283 | print ('Results in ', lr.res_dir) 284 | 285 | if __name__ == '__main__': 286 | sys.exit(main()) 287 | --------------------------------------------------------------------------------