Gonde94
Gonde94

Reputation: 57

My UNet model produces an all gray picture

I'm training a UNet model. It is supposed to take in an apparent resistivity image, shaped like a V and output a prediction for what is underneath the surface.
Here is what I'm getting:
Incorrect prediction

As you can see, for some reason the model is not predicting any blacks or white, just a canvas of gray. The array is something like this:

array([[[0.48203585],
        [0.4981977 ],
        [0.49600083],
        ...,
        [0.49739254],
        [0.49782765],
        [0.48060176]],

       [[0.46797225],
        [0.49758294],
        [0.49503353],
        ...,
        [0.49694082],
        [0.4964686 ],
        [0.4707015 ]],

       [[0.4559476 ],
        [0.49759352],
        [0.49402285],
        ...,
        [0.49774444],
        [0.4961012 ],
        [0.46151617]],

I've made sure my images are all in the correct order, scaled, and everything, but I must be making a mistake somewhere. I've tried binary_crossentropy and mean_squared_error as loss functions. For some reason, the model trains for around 4 epochs before early stopping - what does this mean? Any help would be appreciated. Here is my code:

#Get images
im_width = 128
im_height = 128

train_ids = next(os.walk(dirr+'input/train/'))[2]
tar_ids = next(os.walk(dirr+'target/train/'))[2]

train_ids = sorted(train_ids)
tar_ids = sorted(tar_ids)

X = np.zeros((len(train_ids), im_width, im_height, 1), dtype=np.float32)
y = np.zeros((len(train_ids), im_height, im_width, 1), dtype=np.float32)

for n, id_ in tqdm(enumerate(train_ids), total=len(train_ids)):
    #Load images
    img = load_img(dirr + "input/train/" + id_)
    x_img = img_to_array(img)[:,:,1]
    x_img = resize(x_img, (128, 128, 1), mode='constant', preserve_range=True)
    
    #Normalise images
    X[n] = x_img/255.0
    
for n, id_ in tqdm(enumerate(tar_ids), total=len(tar_ids)):
    #Load targets
    mask = img_to_array(load_img(dirr + "target/train/" + id_))[:,:,1]
    mask = resize(mask, (128, 128, 1), mode='constant', preserve_range=True)
    
    #Normalise images
    y[n] = mask/255.0

Simple UNET for semantic segmentation

def build_model(input_layer, start_neurons):
    # 128 -> 64
    conv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(input_layer)
    conv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(conv1)
    pool1 = MaxPooling2D((2, 2))(conv1)
    pool1 = Dropout(0.25)(pool1)

    # 64 -> 32
    conv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(pool1)
    conv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(conv2)
    pool2 = MaxPooling2D((2, 2))(conv2)
    pool2 = Dropout(0.5)(pool2)

    # 32 -> 16
    conv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(pool2)
    conv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(conv3)
    pool3 = MaxPooling2D((2, 2))(conv3)
    pool3 = Dropout(0.5)(pool3)

    # 16 -> 8
    conv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(pool3)
    conv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(conv4)
    pool4 = MaxPooling2D((2, 2))(conv4)
    pool4 = Dropout(0.5)(pool4)

    # Middle
    convm = Conv2D(start_neurons * 16, (3, 3), activation="relu", padding="same")(pool4)
    convm = Conv2D(start_neurons * 16, (3, 3), activation="relu", padding="same")(convm)

    # 8 -> 16
    deconv4 = Conv2DTranspose(start_neurons * 8, (3, 3), strides=(2, 2), padding="same")(convm)
    uconv4 = concatenate([deconv4, conv4])
    uconv4 = Dropout(0.5)(uconv4)
    uconv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(uconv4)
    uconv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(uconv4)

    # 16 -> 32
    deconv3 = Conv2DTranspose(start_neurons * 4, (3, 3), strides=(2, 2), padding="same")(uconv4)
    uconv3 = concatenate([deconv3, conv3])
    uconv3 = Dropout(0.5)(uconv3)
    uconv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(uconv3)
    uconv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(uconv3)

    # 32 -> 64
    deconv2 = Conv2DTranspose(start_neurons * 2, (3, 3), strides=(2, 2), padding="same")(uconv3)
    uconv2 = concatenate([deconv2, conv2])
    uconv2 = Dropout(0.5)(uconv2)
    uconv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(uconv2)
    uconv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(uconv2)

    # 64 -> 128
    deconv1 = Conv2DTranspose(start_neurons * 1, (3, 3), strides=(2, 2), padding="same")(uconv2)
    uconv1 = concatenate([deconv1, conv1])
    uconv1 = Dropout(0.5)(uconv1)
    uconv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(uconv1)
    uconv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(uconv1)

    #uconv1 = Dropout(0.5)(uconv1)
    output_layer = Conv2D(1, (1,1), padding="same", activation="sigmoid")(uconv1)
    
    return output_layer

input_layer = Input((128, 128, 1))
output_layer = build_model(input_layer, 16)

model = Model(input_layer, output_layer)
model.compile(loss='binary_crossentropy', optimizer='adam', metrics=[jaccard_distance_loss, keras.metrics.RootMeanSquaredError()])
model.summary()

#Train the model

early_stopping = EarlyStopping(patience=2, verbose=1)
model_checkpoint = ModelCheckpoint("FindObjectsMeanIou.h5", save_best_only=True, verbose=1)

epochs = 15
batch_size = 32

history = model.fit(X, y, validation_data=[X_val, y_val], epochs=epochs, batch_size=batch_size,
                   callbacks=[early_stopping, model_checkpoint])

model = load_model(...)

preds_train = model.predict(X, verbose=1)
preds_valid = model.predict(X_val, verbose=1)

#Maybe my visualisations are wrong?

fig, (ax0, ax1, ax2) = plt.subplots(1, 3, figsize=(20, 10))

norm = plt.Normalize(vmin=10, vmax=1200, clip=True)

a = ax0.imshow(np.dstack((X_val[350], X_val[350], X_val[350])), cmap='gray', norm=norm)
ax0.set_title('Apparent Resistivity')
ax0.axis('off')
ab = fig.colorbar(a, ax=ax0, shrink=0.4)
ab.set_label('Resistivity in Ohms')

tmp = np.squeeze(y_val[350]).astype(np.float32)
b = ax1.imshow(np.dstack((tmp,tmp,tmp)), cmap='gray', norm=norm)
ax1.set_title('True Model')
ax1.axis('off')
bb = fig.colorbar(b, ax=ax1, fraction=0.035, pad=0.04)
bb.set_label('Resistivity in Ohms')

tmp = np.squeeze(preds_val[350]).astype(np.float32)
c = ax2.imshow(np.dstack((tmp,tmp,tmp)), cmap='gray', norm=norm)
ax2.set_title('Predicted Model')
ax2.axis('off')
cb = fig.colorbar(c, ax=ax2, fraction=0.035, pad=0.04)

Upvotes: 1

Views: 675

Answers (1)

Azhan Mohammed
Azhan Mohammed

Reputation: 424

Found this implementation of Weighted Binary Cross Entropy Loss from GitHub

import keras

def weighted_bincrossentropy(true, pred, weight_zero = 0.25, weight_one = 1):
    """
    Calculates weighted binary cross entropy. The weights are fixed.
        
    This can be useful for unbalanced catagories.
    
    Adjust the weights here depending on what is required.
    
    For example if there are 10x as many positive classes as negative classes,
        if you adjust weight_zero = 1.0, weight_one = 0.1, then false positives 
        will be penalize 10 times as much as false negatives.
    """
  
    # calculate the binary cross entropy
    bin_crossentropy = keras.backend.binary_crossentropy(true, pred)
    
    # apply the weights
    weights = true * weight_one + (1. - true) * weight_zero
    weighted_bin_crossentropy = weights * bin_crossentropy 

    return keras.backend.mean(weighted_bin_crossentropy)

To get the weights zero and weights one, take any N random images, calculate the number of white and black pixels in the image, and get the ratio. Worked very well for me.

Upvotes: 1

Related Questions