title image segmentation with a unetlike architecture fchollethttps twitter comfchollet date created 20190320 last modified 20200420 description image segmentation model trained from scratch on the oxford pets dataset accelerator gpu download the data shell wget https www robots ox ac ukvggdatapetsdataimages tar gz wget https www robots ox ac ukvggdatapetsdataannotations tar gz curl o https thor robots ox ac ukdatasetspetsimages tar gz curl o https thor robots ox ac ukdatasetspetsannotations tar gz tar xf images tar gz tar xf annotations tar gz prepare paths of input images and target segmentation masks what does one input image and corresponding segmentation mask look like display input image 7 display autocontrast version of corresponding target perpixel categories prepare dataset to load vectorize batches of data returns a tf dataset def loadimgmasksinputimgpath targetimgpath inputimg tfio readfileinputimgpath inputimg tfio decodepnginputimg channels3 inputimg tfimage resizeinputimg imgsize inputimg tfimage convertimagedtypeinputimg float32 targetimg tfio readfiletargetimgpath targetimg tfio decodepngtargetimg channels1 targetimg tfimage resizetargetimg imgsize methodnearest targetimg tfimage convertimagedtypetargetimg uint8 ground truth labels are 1 2 3 subtract one to make them 0 1 2 targetimg 1 return inputimg targetimg for faster debugging limit the size of data if maxdatasetlen inputimgpaths inputimgpaths maxdatasetlen targetimgpaths targetimgpaths maxdatasetlen dataset tfdata dataset fromtensorslices inputimgpaths targetimgpaths dataset dataset maploadimgmasks numparallelcallstfdata autotune return dataset batchbatchsize prepare unet xceptionstyle model from keras import layers def getmodelimgsize numclasses inputs keras inputshapeimgsize 3 first half of the network downsampling inputs entry block x layers conv2d32 3 strides2 paddingsameinputs x layers batchnormalizationx x layers activationrelux previousblockactivation x set aside residual blocks 1 2 3 are identical apart from the feature depth for filters in 64 128 256 x layers activationrelux x layers separableconv2dfilters 3 paddingsamex x layers batchnormalizationx x layers activationrelux x layers separableconv2dfilters 3 paddingsamex x layers batchnormalizationx x layers maxpooling2d3 strides2 paddingsamex project residual residual layers conv2dfilters 1 strides2 paddingsame previousblockactivation x layers addx residual add back residual previousblockactivation x set aside next residual second half of the network upsampling inputs for filters in 256 128 64 32 x layers activationrelux x layers conv2dtransposefilters 3 paddingsamex x layers batchnormalizationx x layers activationrelux x layers conv2dtransposefilters 3 paddingsamex x layers batchnormalizationx x layers upsampling2d2x project residual residual layers upsampling2d2previousblockactivation residual layers conv2dfilters 1 paddingsameresidual x layers addx residual add back residual previousblockactivation x set aside next residual add a perpixel classification layer outputs layers conv2d numclasses 3 activationsoftmax paddingsame x define the model model keras modelinputs outputs return model build model model getmodelimgsize numclasses model summary set aside a validation split import random split our img paths into a training and a validation set valsamples 1000 random random1337 shuffleinputimgpaths random random1337 shuffletargetimgpaths traininputimgpaths inputimgpaths valsamples traintargetimgpaths targetimgpaths valsamples valinputimgpaths inputimgpathsvalsamples valtargetimgpaths targetimgpathsvalsamples instantiate dataset for each split limit input files in maxdatasetlen for faster epoch training time remove the maxdatasetlen arg when running with full dataset traindataset getdataset batchsize imgsize traininputimgpaths traintargetimgpaths maxdatasetlen1000 validdataset getdataset batchsize imgsize valinputimgpaths valtargetimgpaths train the model configure the model for training we use the sparse version of categoricalcrossentropy because our target data is integers model compileoptimizerrmsprop losssparsecategoricalcrossentropy callbacks keras callbacks modelcheckpoint oxfordsegmentation keras savebestonlytrue train the model doing validation at the end of each epoch epochs 15 model fit traindataset epochsepochs validationdatavaliddataset callbackscallbacks visualize predictions generate predictions for all images in the validation set valdataset getdataset batchsize imgsize valinputimgpaths valtargetimgpaths valpreds model predictvaldataset def displaymaski display results for validation image 10 display input image display groundtruth target mask display mask predicted by our model download the data shell wget https www robots ox ac uk vgg data pets data images tar gz wget https www robots ox ac uk vgg data pets data annotations tar gz curl o https thor robots ox ac uk datasets pets images tar gz curl o https thor robots ox ac uk datasets pets annotations tar gz tar xf images tar gz tar xf annotations tar gz prepare paths of input images and target segmentation masks what does one input image and corresponding segmentation mask look like display input image 7 display auto contrast version of corresponding target per pixel categories prepare dataset to load vectorize batches of data returns a tf dataset ground truth labels are 1 2 3 subtract one to make them 0 1 2 for faster debugging limit the size of data prepare u net xception style model first half of the network downsampling inputs entry block set aside residual blocks 1 2 3 are identical apart from the feature depth project residual add back residual set aside next residual second half of the network upsampling inputs project residual add back residual set aside next residual add a per pixel classification layer define the model build model set aside a validation split split our img paths into a training and a validation set instantiate dataset for each split limit input files in max_dataset_len for faster epoch training time remove the max_dataset_len arg when running with full dataset train the model configure the model for training we use the sparse version of categorical_crossentropy because our target data is integers train the model doing validation at the end of each epoch visualize predictions generate predictions for all images in the validation set quick utility to display a model s prediction display results for validation image 10 display input image display ground truth target mask display mask predicted by our model note that the model only sees inputs at 150x150
import os

input_dir = "images/"
target_dir = "annotations/trimaps/"
img_size = (160, 160)
num_classes = 3
batch_size = 32

input_img_paths = sorted(
    [
        os.path.join(input_dir, fname)
        for fname in os.listdir(input_dir)
        if fname.endswith(".jpg")
    ]
)
target_img_paths = sorted(
    [
        os.path.join(target_dir, fname)
        for fname in os.listdir(target_dir)
        if fname.endswith(".png") and not fname.startswith(".")
    ]
)

print("Number of samples:", len(input_img_paths))

for input_path, target_path in zip(input_img_paths[:10], target_img_paths[:10]):
    print(input_path, "|", target_path)



from IPython.display import Image, display
from keras.utils import load_img
from PIL import ImageOps


display(Image(filename=input_img_paths[9]))


img = ImageOps.autocontrast(load_img(target_img_paths[9]))
display(img)



import keras
import numpy as np
from tensorflow import data as tf_data
from tensorflow import image as tf_image
from tensorflow import io as tf_io


def get_dataset(
    batch_size,
    img_size,
    input_img_paths,
    target_img_paths,
    max_dataset_len=None,
):
    

    def load_img_masks(input_img_path, target_img_path):
        input_img = tf_io.read_file(input_img_path)
        input_img = tf_io.decode_png(input_img, channels=3)
        input_img = tf_image.resize(input_img, img_size)
        input_img = tf_image.convert_image_dtype(input_img, "float32")

        target_img = tf_io.read_file(target_img_path)
        target_img = tf_io.decode_png(target_img, channels=1)
        target_img = tf_image.resize(target_img, img_size, method="nearest")
        target_img = tf_image.convert_image_dtype(target_img, "uint8")

        
        target_img -= 1
        return input_img, target_img

    
    if max_dataset_len:
        input_img_paths = input_img_paths[:max_dataset_len]
        target_img_paths = target_img_paths[:max_dataset_len]
    dataset = tf_data.Dataset.from_tensor_slices(
        (input_img_paths, target_img_paths)
    )
    dataset = dataset.map(load_img_masks, num_parallel_calls=tf_data.AUTOTUNE)
    return dataset.batch(batch_size)




from keras import layers


def get_model(img_size, num_classes):
    inputs = keras.Input(shape=img_size + (3,))

    

    
    x = layers.Conv2D(32, 3, strides=2, padding="same")(inputs)
    x = layers.BatchNormalization()(x)
    x = layers.Activation("relu")(x)

    previous_block_activation = x  

    
    for filters in [64, 128, 256]:
        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.SeparableConv2D(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.MaxPooling2D(3, strides=2, padding="same")(x)

        
        residual = layers.Conv2D(filters, 1, strides=2, padding="same")(
            previous_block_activation
        )
        x = layers.add([x, residual])  
        previous_block_activation = x  

    

    for filters in [256, 128, 64, 32]:
        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.Activation("relu")(x)
        x = layers.Conv2DTranspose(filters, 3, padding="same")(x)
        x = layers.BatchNormalization()(x)

        x = layers.UpSampling2D(2)(x)

        
        residual = layers.UpSampling2D(2)(previous_block_activation)
        residual = layers.Conv2D(filters, 1, padding="same")(residual)
        x = layers.add([x, residual])  
        previous_block_activation = x  

    
    outputs = layers.Conv2D(
        num_classes, 3, activation="softmax", padding="same"
    )(x)

    
    model = keras.Model(inputs, outputs)
    return model



model = get_model(img_size, num_classes)
model.summary()



import random


val_samples = 1000
random.Random(1337).shuffle(input_img_paths)
random.Random(1337).shuffle(target_img_paths)
train_input_img_paths = input_img_paths[:-val_samples]
train_target_img_paths = target_img_paths[:-val_samples]
val_input_img_paths = input_img_paths[-val_samples:]
val_target_img_paths = target_img_paths[-val_samples:]




train_dataset = get_dataset(
    batch_size,
    img_size,
    train_input_img_paths,
    train_target_img_paths,
    max_dataset_len=1000,
)
valid_dataset = get_dataset(
    batch_size, img_size, val_input_img_paths, val_target_img_paths
)






model.compile(optimizer="rmsprop", loss="sparse_categorical_crossentropy")

callbacks = [
    keras.callbacks.ModelCheckpoint(
        "oxford_segmentation.keras", save_best_only=True
    )
]


epochs = 15
model.fit(
    train_dataset,
    epochs=epochs,
    validation_data=valid_dataset,
    callbacks=callbacks,
)





val_dataset = get_dataset(
    batch_size, img_size, val_input_img_paths, val_target_img_paths
)
val_preds = model.predict(val_dataset)


def display_mask(i):
    
    mask = np.argmax(val_preds[i], axis=-1)
    mask = np.expand_dims(mask, axis=-1)
    img = ImageOps.autocontrast(keras.utils.array_to_img(mask))
    display(img)



i = 10


display(Image(filename=val_input_img_paths[i]))


img = ImageOps.autocontrast(load_img(val_target_img_paths[i]))
display(img)


display_mask(i)