Keras: Image segmentation using grayscale masks and ImageDataGenerator class

Question

I am currently trying to implement a convolutional network using Keras 2.1.6 (with TensorFlow as backend) and its ImageDataGenerator to segment an image using a grayscale mask. I try to use an image as input, and a mask as label. Due to a low amount of training images, and memory constraints I utilize the ImageDataGenerator class provided in Keras.

However I get this error, after changing the values provided in the Keras example to the ones described later:

File "C:\Users\XXX\Anaconda3\lib\site-packages\keras\engine\training.py", line 2223, in fit_generator
batch_size = x.shape[0]
AttributeError: 'tuple' object has no attribute 'shape'

Which, as far as I know, happens because the generator does generate a tuple, and not an array. This first happened after I changed following parameters from the standard values provided in the Keras example to the following: color_mode='grayscale' for all mask generators, and class_mode='input' due to this being recommended for autoencoders.

The Keras example can be found in here.

The dataset I am using consists of 100 images (jpg) and 100 corresponding grayscale masks (png) and can be downloaded at this link

The architecture I wanted to implement is an autoencoder/U-Net based network and it is shown in the provided code:

from keras.preprocessing import image
from keras.models import Model
from keras import optimizers 
from keras.preprocessing.image import ImageDataGenerator
from keras.layers import Input, Conv2D, MaxPooling2D, UpSampling2D
from keras import initializers

image_path = 
mask_path = 
valid_image_path = 
valid_mask_path = 

img_size=160
batchsize=10
samplesize = 60
steps = samplesize / batchsize

train_datagen = image.ImageDataGenerator(shear_range=0.2,
                                        zoom_range=0.2,
                                        horizontal_flip=True)

data_gen_args = dict(rotation_range=90,
                     width_shift_range=0.1,
                     height_shift_range=0.1,
                     zoom_range=0.2)

image_datagen = ImageDataGenerator(**data_gen_args)
mask_datagen = ImageDataGenerator(**data_gen_args)

seed = 1

image_generator = image_datagen.flow_from_directory(
    image_path,
    target_size=(img_size, img_size),
    class_mode='input',
    batch_size = batchsize,
    seed=seed)

mask_generator = mask_datagen.flow_from_directory(
    mask_path,
    target_size=(img_size, img_size),
    class_mode='input',
    color_mode = 'grayscale',
    batch_size = batchsize,
    seed=seed)

vimage_generator = image_datagen.flow_from_directory(
    valid_image_path,
    target_size=(img_size, img_size),
    class_mode='input',
    batch_size = batchsize,
    seed=seed)

vmask_generator = mask_datagen.flow_from_directory(
    valid_mask_path,
    target_size=(img_size, img_size),
    class_mode='input',
    color_mode = 'grayscale',
    batch_size = batchsize,
    seed=seed)

#Model 
input_img = Input(shape=(img_size,img_size,3))
c11 = Conv2D(16, (3, 3), activation='relu', padding='same', kernel_initializer=initializers.random_normal(stddev=0.01))(input_img)
mp1 = MaxPooling2D((2, 2), padding='same')(c11)
c21 = Conv2D(16, (3, 3), activation='relu', padding='same', kernel_initializer=initializers.random_normal(stddev=0.01))(mp1)
mp2 = MaxPooling2D((2, 2), padding='same')(c21)
c31 = Conv2D(32, (3, 3), activation='relu', padding='same', kernel_initializer=initializers.random_normal(stddev=0.01))(mp2)
encoded = MaxPooling2D((5, 5), padding='same')(c31)

c12 = Conv2D(32, (3, 3), activation='relu', padding='same', kernel_initializer=initializers.random_normal(stddev=0.01))(encoded)
us12 = UpSampling2D((5,5))(c12)
c22 = Conv2D(16, (3, 3), activation='relu', padding='same', kernel_initializer=initializers.random_normal(stddev=0.01))(us12)
us22 = UpSampling2D((2, 2))(c22)
c32 = Conv2D(16, (3, 3), activation='relu', padding='same', kernel_initializer=initializers.random_normal(stddev=0.01))(us22)
us32 = UpSampling2D((2, 2))(c32)
decoded = Conv2D(1, (3, 3), activation='softmax', padding='same')(us32)

model = Model(input_img, decoded)

model.compile(loss="mean_squared_error", optimizer=optimizers.Adam(),metrics=["accuracy"])
#model.summary()

#Generators, tr: training, v: validation
trgen = zip(image_generator,mask_generator)
vgen = zip(vimage_generator,vmask_generator)

model.fit_generator(
    trgen,
    steps_per_epoch= steps,
    epochs=5,
    validation_data = vgen,
    validation_steps=10)

Belkisbelknap

Answer 1

Here is a better version of Unet, you can use this code

def conv_block(tensor, nfilters, size=3, padding='same', initializer="he_normal"):
    x = Conv2D(filters=nfilters, kernel_size=(size, size), padding=padding, kernel_initializer=initializer)(tensor)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    x = Conv2D(filters=nfilters, kernel_size=(size, size), padding=padding, kernel_initializer=initializer)(x)
    x = BatchNormalization()(x)
    x = Activation("relu")(x)
    return x


def deconv_block(tensor, residual, nfilters, size=3, padding='same', strides=(2, 2)):
    y = Conv2DTranspose(nfilters, kernel_size=(size, size), strides=strides, padding=padding)(tensor)
    y = concatenate([y, residual], axis=3)
    y = conv_block(y, nfilters)
    return y


def Unet(img_height, img_width, nclasses=3, filters=64):
# down
    input_layer = Input(shape=(img_height, img_width, 3), name='image_input')
    conv1 = conv_block(input_layer, nfilters=filters)
    conv1_out = MaxPooling2D(pool_size=(2, 2))(conv1)
    conv2 = conv_block(conv1_out, nfilters=filters*2)
    conv2_out = MaxPooling2D(pool_size=(2, 2))(conv2)
    conv3 = conv_block(conv2_out, nfilters=filters*4)
    conv3_out = MaxPooling2D(pool_size=(2, 2))(conv3)
    conv4 = conv_block(conv3_out, nfilters=filters*8)
    conv4_out = MaxPooling2D(pool_size=(2, 2))(conv4)
    conv4_out = Dropout(0.5)(conv4_out)
    conv5 = conv_block(conv4_out, nfilters=filters*16)
    conv5 = Dropout(0.5)(conv5)
# up
    deconv6 = deconv_block(conv5, residual=conv4, nfilters=filters*8)
    deconv6 = Dropout(0.5)(deconv6)
    deconv7 = deconv_block(deconv6, residual=conv3, nfilters=filters*4)
    deconv7 = Dropout(0.5)(deconv7) 
    deconv8 = deconv_block(deconv7, residual=conv2, nfilters=filters*2)
    deconv9 = deconv_block(deconv8, residual=conv1, nfilters=filters)
# output
    output_layer = Conv2D(filters=nclasses, kernel_size=(1, 1))(deconv9)
    output_layer = BatchNormalization()(output_layer)
    output_layer = Activation('softmax')(output_layer)

    model = Model(inputs=input_layer, outputs=output_layer, name='Unet')
    return model

Note if you have only two classes ie nclasses=2, you need to change

output_layer = Conv2D(filters=nclasses, kernel_size=(1, 1))(deconv9)
output_layer = BatchNormalization()(output_layer)
output_layer = Activation('softmax')(output_layer)

to

output_layer = Conv2D(filters=2, kernel_size=(1, 1))(deconv9)
output_layer = BatchNormalization()(output_layer)
output_layer = Activation('sigmoid')(output_layer)

Now for the data generators, you can use the builtin ImageDataGenerator class here is the code from Keras docs

# we create two instances with the same arguments
data_gen_args = dict(featurewise_center=True,
                     featurewise_std_normalization=True,
                     rotation_range=90,
                     width_shift_range=0.1,
                     height_shift_range=0.1,
                     zoom_range=0.2)
image_datagen = ImageDataGenerator(**data_gen_args)
mask_datagen = ImageDataGenerator(**data_gen_args)

# Provide the same seed and keyword arguments to the fit and flow methods
seed = 1
image_datagen.fit(images, augment=True, seed=seed)
mask_datagen.fit(masks, augment=True, seed=seed)

image_generator = image_datagen.flow_from_directory(
    'data/images',
    class_mode=None,
    seed=seed)

mask_generator = mask_datagen.flow_from_directory(
    'data/masks',
    class_mode=None,
    seed=seed)

# combine generators into one which yields image and masks
train_generator = zip(image_generator, mask_generator)

model.fit_generator(
    train_generator,
    steps_per_epoch=2000,
    epochs=50)

Another way to go is implement your own generator by extending the Sequence class from Keras

class seg_gen(Sequence):
    def __init__(self, x_set, y_set, batch_size, image_dir, mask_dir):
        self.x, self.y = x_set, y_set
        self.batch_size = batch_size
        self.samples = len(self.x)
        self.image_dir = image_dir
        self.mask_dir = mask_dir

    def __len__(self):
        return int(np.ceil(len(self.x) / float(self.batch_size)))

    def __getitem__(self, idx):
        idx = np.random.randint(0, self.samples, batch_size)
        batch_x, batch_y = [], []
        drawn = 0
        for i in idx:
            _image = image.img_to_array(image.load_img(f'{self.image_dir}/{self.x[i]}', target_size=(img_height, img_width)))/255.   
            mask = image.img_to_array(image.load_img(f'{self.mask_dir}/{self.y[i]}', grayscale=True, target_size=(img_height, img_width)))
#             mask = np.resize(mask,(img_height*img_width, classes))
            batch_y.append(mask)
            batch_x.append(_image)
        return np.array(batch_x), np.array(batch_y)

Here is a sample code to train the model

unet = Unet(256, 256, nclasses=66, filters=64)
print(unet.output_shape)
p_unet = multi_gpu_model(unet, 4)
p_unet.load_weights('models-dr/top_weights.h5')
p_unet.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
tb = TensorBoard(log_dir='logs', write_graph=True)
mc = ModelCheckpoint(mode='max', filepath='models-dr/top_weights.h5', monitor='acc', save_best_only='True', save_weights_only='True', verbose=1)
es = EarlyStopping(mode='max', monitor='acc', patience=6, verbose=1)
callbacks = [tb, mc, es]
train_gen = seg_gen(image_list, mask_list, batch_size)


p_unet.fit_generator(train_gen, steps_per_epoch=steps, epochs=13, callbacks=callbacks, workers=8)

I got good results when i had only 2 classes by using dice loss, here is the code for it

def dice_coeff(y_true, y_pred):
    smooth = 1.
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    score = (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
    return score

def dice_loss(y_true, y_pred):
    loss = 1 - dice_coeff(y_true, y_pred)
    return loss

Punctate

When you use data augmentation on your y data (masks), won't this change the values (e.g., for rotation or shift it may use interpolation to calculate pixel values; other methods may directly change color values). Isn't this a problem, since the values in the mask images correspond to classes? If 0 = house and 1 = dog and you rotate the image, the outcome could be some pixels with value 0.5, or is there something to prevent this? – Intercut

this is true, but you can take care of this by setting fill_mode=nearest in ImageDataGenerator args. The deafult is already set to nearest. – Punctate

@SrihariHumbarwadi from where this 'images' is coming in [image_datagen.fit(images, augment=True, seed=seed)] ? – Santalaceous

if I used my own generator by extending the Sequence class, how to use ImageDataGenerator class in this case. – Damico

Answer 2

What you are trying to build is an image segmentation model and not an autoencoder. Therefore, since you have separate generators for the images and the labels (i.e. masks), you need to set the class_mode argument to None to prevent generator from producing any labels arrays.

Further, you need to change the activation function of last layer from softmax to sigmoid, otherwise since the softmax normalizes the sum of its input elements to 1, the output would be all ones. You can also use binary_crossentropy for the loss function as well.

Cutlerr

Thanks, this lets me at least run my code. However, the training accuracy decreases during training now, i.e. instead of overfitting, the model begins to underfit more udner training. – Belkisbelknap

@Belkisbelknap One reason might be that the input images are not normalized. Pass rescale=1/255.0 when constructing image_datagen. Further, if the masks values have been stored as 0 and 255 then you also need to rescale them the same way to make them 0 and 1. However, if they are already 0 and 1, don't use rescaling for them. – Cutlerr