example/fcn-xs/image_segmentaion.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

"""
This module encapsulates running image segmentation model for inference.

Example usage:
    $ python image_segmentaion.py --input <your JPG image path>
"""

import argparse
import os
import numpy as np
import mxnet as mx
from PIL import Image

def make_file_extension_assertion(extension):
    """Function factory for file extension argparse assertion
        Args:
            extension (string): the file extension to assert

        Returns:
            string: the supplied extension, if assertion is successful.

    """
    def file_extension_assertion(file_path):
        base, ext = os.path.splitext(file_path)
        if ext.lower() != extension:
            raise argparse.ArgumentTypeError('File must have ' + extension + ' extension')
        return file_path
    return file_extension_assertion

def get_palette(num_colors=256):
    """generates the colormap for visualizing the segmentation mask
            Args:
                num_colors (int): the number of colors to generate in the output palette

            Returns:
                string: the supplied extension, if assertion is successful.

    """
    pallete = [0]*(num_colors*3)
    for j in range(0, num_colors):
        lab = j
        pallete[j*3+0] = 0
        pallete[j*3+1] = 0
        pallete[j*3+2] = 0
        i = 0
        while (lab > 0):
            pallete[j*3+0] |= (((lab >> 0) & 1) << (7-i))
            pallete[j*3+1] |= (((lab >> 1) & 1) << (7-i))
            pallete[j*3+2] |= (((lab >> 2) & 1) << (7-i))
            i = i + 1
            lab >>= 3
    return pallete

def get_data(img_path):
    """get the (1, 3, h, w) np.array data for the supplied image
                Args:
                    img_path (string): the input image path

                Returns:
                    np.array: image data in a (1, 3, h, w) shape

    """
    mean = np.array([123.68, 116.779, 103.939])  # (R,G,B)
    img = Image.open(img_path)
    img = np.array(img, dtype=np.float32)
    reshaped_mean = mean.reshape(1, 1, 3)
    img = img - reshaped_mean
    img = np.swapaxes(img, 0, 2)
    img = np.swapaxes(img, 1, 2)
    img = np.expand_dims(img, axis=0)
    return img

def main():
    """Module main execution"""
    # Initialization variables - update to change your model and execution context
    model_prefix = "FCN8s_VGG16"
    epoch = 19

    # By default, MXNet will run on the CPU. Change to ctx = mx.gpu() to run on GPU.
    ctx = mx.cpu()

    fcnxs, fcnxs_args, fcnxs_auxs = mx.model.load_checkpoint(model_prefix, epoch)
    fcnxs_args["data"] = mx.nd.array(get_data(args.input), ctx)
    data_shape = fcnxs_args["data"].shape
    label_shape = (1, data_shape[2]*data_shape[3])
    fcnxs_args["softmax_label"] = mx.nd.empty(label_shape, ctx)
    exector = fcnxs.bind(ctx, fcnxs_args, args_grad=None, grad_req="null", aux_states=fcnxs_args)
    exector.forward(is_train=False)
    output = exector.outputs[0]
    out_img = np.uint8(np.squeeze(output.asnumpy().argmax(axis=1)))
    out_img = Image.fromarray(out_img)
    out_img.putpalette(get_palette())
    out_img.save(args.output)

if __name__ == "__main__":
    # Handle command line arguments
    parser = argparse.ArgumentParser(description='Run VGG16-FCN-8s to segment an input image')
    parser.add_argument('--input',
                        required=True,
                        type=make_file_extension_assertion('.jpg'),
                        help='The segmentation input JPG image')
    parser.add_argument('--output',
                        default='segmented.png',
                        type=make_file_extension_assertion('.png'),
                        help='The segmentation putput PNG image')
    args = parser.parse_args()
    main()