Search code examples
kerascomputer-visionimage-segmentationmobilenet

How can I feed the output from last layer of mobilenet to a Unet model


I am trying to build an image segmentation model with a Keras mobilenet model pre-trained on imagenet dataset. How ever to train the model further, I want to add the U-net layers to the existing model and only train the layers of u-net architecture with mobilenet model helping as a backbone.

Problem: The last layer of mobilenet model is of dimensions (7x7x1024), which is a RelU layer, I wish want to re-shape this to (256x256x3) which can be understood by the U-net input layer.


Solution

  • not the last layer, but creating a unet on mobilenet can be done using the below code:

    ALPHA = 1 # Width hyper parameter for MobileNet (0.25, 0.5, 0.75, 1.0). Higher width means more accurate but slower
    
    IMAGE_HEIGHT = 224
    IMAGE_WIDTH = 224
    
    HEIGHT_CELLS = 28
    WIDTH_CELLS = 28
    
    def create_model(trainable=True):
        model = MobileNet(input_shape=(IMAGE_HEIGHT, IMAGE_WIDTH, 3), include_top=False, alpha=ALPHA, weights="imagenet")
    
        block0 = model.get_layer("conv_pw_1_relu").output 
        block = model.get_layer("conv_pw_1_relu").output
        block1 = model.get_layer("conv_pw_3_relu").output
        block2 = model.get_layer("conv_pw_5_relu").output
        block3 = model.get_layer("conv_pw_11_relu").output
        block4 = model.get_layer("conv_pw_13_relu").output
    
        x = Concatenate()([UpSampling2D()(block4), block3])
        x = Concatenate()([UpSampling2D()(x), block2])
        x = Concatenate()([UpSampling2D()(x), block1])
        x = Concatenate()([UpSampling2D()(x), block])
     #   x = Concatenate()([UpSampling2D()(x), block0])
        x = UpSampling2D()(x)
        x = Conv2D(1, kernel_size=1, activation="sigmoid")(x)
    
        x = Reshape((IMAGE_HEIGHT, IMAGE_HEIGHT))(x)
    
        return Model(inputs=model.input, outputs=x)