├── .gitignore ├── README.md ├── models └── unet_1 │ ├── __init__.py │ └── sessions │ └── session_1 │ └── weights.hd5 ├── test.py ├── train.py └── utility.py /.gitignore: -------------------------------------------------------------------------------- 1 | .idea 2 | __pycache__/ 3 | *.png 4 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Segmentation-Experiment 2 | Using Keras (U-Net architecture) to segment shapes on noise. 3 | 4 | **General**: 5 | 6 | I created this experiment to toy around with some convolutional networks for image segmentation, mostly U-Net architecture (https://lmb.informatik.uni-freiburg.de/people/ronneber/u-net/). 7 | For this I create a noise-image and at a random location within that noise-image I put a circle of random size. 8 | This circle can be visually different from the noise in different ways: taking constant values or a little bit stronger noise. 9 | 10 | The reason I put this Code online was mostly to discuss possible interpretation of the interesting observation: Training the network (~13 million parameters) 11 | on only 1 training sample (that means: only 1 image with a circle of random size at a random location) is already sufficient to make 12 | the network converge on a validation/test set (i.e. on noise-images with circles of other sizes and other locations the network was not trained on). 13 | 14 | This result surprised me because due to the size of the network I expected the network to simply 'memorize' (i.e. overfit) 15 | on the one training sample that I provided. Instead it seems to find the simplest hypothesis (which is not memorization) 16 | that is sufficient to fit the training sample, and this hypothesis seems to be the one we actually seek. 17 | 18 | The following image shows the training sample that is produced by 'train.py' and the one image I used to train the network with: 19 | 20 | ![Training Sample](http://i.imgur.com/i9OQfNb.png) 21 | 22 | The network is then trained to segment the circle. 23 | 24 | 25 | **Code**: 26 | 27 | Just run train.py to load the u-net model and train the model. The function 'create_data' will create the training and validation 28 | data set and save the images to the 'output' directory (I like visualizing my data). 'X_' and 'Y_' denote the noise-images and the ground-truth mask respectively. 29 | 30 | The test.py file tests the trained network on some newly created data by showing ground-truth and the predicted mask in one image which is saved 31 | in the 'vis' directory. 32 | 33 | 34 | 35 | 36 | 37 | **My current setup**: 38 | - Windows 10 64 bit with an Nvidia 980 TI 39 | - Python 3.6 40 | - opencv 3.2.0 41 | - Tensorflow 1.2.0-rc1 42 | - Keras 2.0.4 43 | -------------------------------------------------------------------------------- /models/unet_1/__init__.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from keras.models import Model 3 | from keras.layers import Activation, Input, concatenate, Conv2D, MaxPooling2D, UpSampling2D, Dropout, Reshape, Permute, Concatenate, concatenate 4 | from keras.optimizers import Adam, SGD 5 | from keras.callbacks import ModelCheckpoint, LearningRateScheduler 6 | from keras import backend as K 7 | from keras.optimizers import Optimizer 8 | 9 | img_rows = 160 10 | img_cols = 160 11 | 12 | smooth = 1. 13 | 14 | def mask_binary_regression_error(y_true, y_pred): 15 | return K.mean(K.log(1 + K.exp(-y_true*y_pred))) 16 | 17 | 18 | def dice_coef(y_true, y_pred): 19 | y_true_f = K.flatten(y_true) 20 | y_pred_f = K.flatten(y_pred) 21 | intersection = K.sum(y_true_f * y_pred_f) 22 | return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth) 23 | 24 | 25 | def dice_coef_loss(y_true, y_pred): 26 | return 1-dice_coef(y_true, y_pred) 27 | 28 | inputs = Input((img_rows, img_cols, 1)) # 160 x 160 29 | conv1 = Conv2D(32, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(inputs) 30 | conv1 = Conv2D(32, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv1) 31 | pool1 = MaxPooling2D(pool_size=(2, 2))(conv1) 32 | # pool1 = Dropout(0.15)(pool1) 33 | 34 | conv2 = Conv2D(64, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(pool1) 35 | conv2 = Conv2D(64, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv2) 36 | pool2 = MaxPooling2D(pool_size=(2, 2))(conv2) 37 | # pool2 = Dropout(0.25)(pool2) 38 | 39 | conv3 = Conv2D(128, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(pool2) 40 | conv3 = Conv2D(128, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv3) 41 | pool3 = MaxPooling2D(pool_size=(2, 2))(conv3) 42 | # pool3 = Dropout(0.4)(pool3) 43 | 44 | conv4 = Conv2D(256, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(pool3) 45 | conv4 = Conv2D(256, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv4) 46 | pool4 = MaxPooling2D(pool_size=(2, 2))(conv4) 47 | # pool4 = Dropout(0.5)(pool4) 48 | 49 | conv5 = Conv2D(512, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(pool4) 50 | conv5 = Conv2D(512, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv5) 51 | pool5 = MaxPooling2D(pool_size=(2, 2))(conv5) # 5x5 52 | # pool5 = Dropout(0.5)(pool5) 53 | 54 | conv6 = Conv2D(512, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(pool5) 55 | conv6 = Conv2D(512, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv6) 56 | 57 | 58 | up7 = concatenate([UpSampling2D(size=(2, 2))(conv6), conv5], axis = 3) 59 | # up7 = Dropout(0.5)(up7) 60 | conv7 = Conv2D(256, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(up7) 61 | conv7 = Conv2D(256, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv7) 62 | 63 | 64 | up8 = concatenate([UpSampling2D(size=(2, 2))(conv7), conv4], axis = 3) 65 | # up8 = Dropout(0.4)(up8) 66 | conv8 = Conv2D(128, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(up8) 67 | conv8 = Conv2D(128, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv8) 68 | 69 | up9 = concatenate([UpSampling2D(size=(2, 2))(conv8), conv3], axis = 3) 70 | # up9 = Dropout(0.25)(up9) 71 | conv9 = Conv2D(64, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(up9) 72 | conv9 = Conv2D(64, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv9) 73 | 74 | up10 = concatenate([UpSampling2D(size=(2, 2))(conv9), conv2], axis = 3) 75 | # up10 = Dropout(0.15)(up10) 76 | conv10 = Conv2D(32, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(up10) 77 | conv10 = Conv2D(32, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv10) 78 | 79 | up11 = concatenate([UpSampling2D(size=(2, 2))(conv10), conv1], axis = 3) 80 | # up11 = Dropout(0.15)(up11) 81 | conv11 = Conv2D(16, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(up11) 82 | conv11 = Conv2D(16, kernel_size = (3,3), strides = 1, activation='relu', padding='same', kernel_initializer = 'he_normal')(conv11) 83 | 84 | conv12 = Conv2D(1, kernel_size = (1,1), strides = 1, activation='sigmoid', kernel_initializer = 'he_normal')(conv11) # n x 2 x 160 x 160 85 | # conv12 = Conv2D(2, 1, 1, kernel_initializer = 'he_normal')(conv11) # n x 2 x 160 x 160 86 | 87 | # r = Reshape((2, 160*160))(conv12) 88 | # r = Permute((2,1))(r) 89 | # r = Activation('softmax')(r) 90 | # r = Activation('sigmoid')(conv12) 91 | 92 | 93 | model = Model(input=inputs, output=conv12) 94 | 95 | options = { 96 | 'session_1': 97 | { 98 | 'loss' : dice_coef_loss, 99 | # 'loss' : 'binary_crossentropy', 100 | 'optimizer': SGD(lr = 0.001), 101 | 'metrics' : [dice_coef] 102 | } 103 | 104 | } 105 | 106 | -------------------------------------------------------------------------------- /models/unet_1/sessions/session_1/weights.hd5: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/mdfwn/Segmentation-Experiment/041f1b1b79463e6fce5af6888858b97298db3c56/models/unet_1/sessions/session_1/weights.hd5 -------------------------------------------------------------------------------- /test.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import os 4 | from utility import create_data 5 | 6 | from imp import load_source 7 | 8 | h,w = 160, 160 # image dimensions 9 | Train = 1 # Number of samples used for training 10 | Val = 100 # Number of samples used for validation during training 11 | np.random.seed(2) 12 | # Take these from train.py 13 | std = 0.09 14 | mean = 0.15 15 | X_test, Y_test = create_data(h, w, 100, size = (15, 40), mode = 'circle') 16 | 17 | # Loading the model: 18 | model_base = 'models' 19 | model_id = 'unet_1' 20 | session_id = 'session_1' 21 | model_path = os.path.join(model_base, model_id) 22 | session_path = os.path.join(model_path, 'sessions', session_id) 23 | session = load_source('session', os.path.join(model_path, '__init__.py')) 24 | Model = session.model 25 | options = session.options 26 | 27 | Model.load_weights( os.path.join(session_path, 'weights.hd5') ) 28 | 29 | for i,x in enumerate(X_test): 30 | Y = Y_test[i] 31 | img = ((x.reshape((h, w)) * std) + mean)*255 32 | img = img.astype('uint8') 33 | img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR) 34 | 35 | R = (Model.predict(x.reshape((1, h, w,1)))*255).reshape((h, w)).astype('uint8') 36 | Y = (Y*255).reshape((h, w)).astype('uint8') 37 | mask = np.zeros((h, w,3), dtype = 'uint8') 38 | mask[:,:,0] = Y 39 | mask[:,:,2] = R 40 | 41 | cv2.imwrite('vis/test_{}.png'.format(i), mask) -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | import os, sys 4 | from utility import create_data 5 | from imp import load_source 6 | 7 | h,w = 160, 160 # image dimensions 8 | Train = 1 # Number of samples used for training 9 | Val = 100 # Number of samples used for validation during training 10 | np.random.seed(1) 11 | 12 | 13 | X_train, Y_train = create_data(h, w, Train, size = (15,40), mode = 'circle') 14 | sys.exit() 15 | # If we set Train = 1, i.e. only 1 training sample, we can speed up training by 16 | # having 100 samples that are all identical 17 | X_train = np.repeat(X_train, 100, 0) 18 | Y_train = np.repeat(Y_train, 100, 0) 19 | 20 | X_val, Y_val = create_data(h, w, Val, size = (15,40), mode = 'circle') 21 | 22 | # Loading the model: 23 | model_base = 'models' 24 | model_id = 'unet_1' 25 | session_id = 'session_1' 26 | model_path = os.path.join(model_base, model_id) 27 | session_path = os.path.join(model_path, 'sessions', session_id) 28 | session = load_source('session', os.path.join(model_path, '__init__.py')) 29 | Model = session.model 30 | options = session.options 31 | 32 | print (Model.summary()) 33 | Model.compile(**options[session_id]) 34 | Model.fit(x = X_train, y = Y_train, nb_epoch = 20, batch_size = 1, verbose = True, validation_data = (X_val, Y_val)) 35 | Model.save_weights( os.path.join(session_path, 'weights.hd5') ) -------------------------------------------------------------------------------- /utility.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | def create_data(h, w, N, size, mode = 'circle'): 5 | ''' 6 | Function to create training samples - N images of size h x w which are 7 | saved in the directory 'output'. 8 | 9 | h,w : height and width (image dimensions) 10 | N : number of images 11 | mode : 'circle' will create circles, 'rectangle' will create rectangles 12 | size: tuple containing lower bound and upper bound for the size of the circle/rectangle 13 | fill : how to fill the circles in the image; 'noise' will fill 14 | ''' 15 | 16 | X = np.zeros((N, h, w), dtype = 'uint8') 17 | Y = np.zeros((N, h, w), dtype = 'uint8') 18 | for i, M in enumerate(Y): 19 | 20 | if mode == 'circle': 21 | r = np.random.randint(size[0], size[1]) 22 | # chose position of circle such that all of the circle will be within the image: 23 | y, x = np.random.randint(2 * r, h - 2 * r), np.random.randint(2 * r, w - 2 * r) 24 | cv2.circle(M, (x,y), r, 255, -1) 25 | 26 | elif mode == 'rectangle': 27 | height, width = np.random.randint(size[0], size[1], 2) 28 | # height = width # Using squares 29 | y = np.random.randint(0, h-height-1) 30 | x = np.random.randint(0, w-width-1) 31 | M[y:y+height, x:x+width] = 255 32 | 33 | Y[i,:,:] = M 34 | 35 | noise1 = (np.random.randint(0, 80, size = (h,w))).astype('uint8') 36 | noise2 = (np.random.randint(0, 80, size = (h,w))).astype('uint8') 37 | 38 | # cv2.circle(noise1, (x,y), r, np.mean(M) ,-1) # Constant value for circles? 39 | noise1 = cv2.add(noise1,(M/255).astype('uint8') * noise2) 40 | 41 | X[i, :, :] = noise1 42 | cv2.imwrite('output/X_{}.png'.format(i+1), noise1) 43 | cv2.imwrite('output/Y_{}.png'.format(i+1), M) 44 | Y = Y.astype('float32').reshape((N,h,w,1)) / 255.0 45 | X = X.astype('float32').reshape((N,h,w,1)) / 255.0 46 | mean, std = np.mean(X), np.std(X) 47 | print (mean, std) 48 | X = (X - mean)/std 49 | 50 | return X, Y --------------------------------------------------------------------------------