nolanliou / mobile-deeplab-v3-plus

Deeplab-V3+ model with MobilenetV2/MobilenetV3 on TensorFlow for mobile deployment.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

I tried to write predictive model code and visualize it, but the model output is not a segmented image

shakey-cuimiao opened this issue · comments

I write test model code, but the model output is not mask, there is no result as shown in the picture below, can you tell me how to get
image

` return colormap[label]

def write_result(imagefile, imgname, outputDir):
#!---- coding:utf-8 ----

Deeplab Demo

import os
import numpy as np
from PIL import Image
import tempfile
from six.moves import urllib
class DeepLabModel(object):
"""
加载 DeepLab 模型;
推断 Inference.
"""
INPUT_TENSOR_NAME = 'Input:0'
OUTPUT_TENSOR_NAME = 'Output:0'
INPUT_SIZE = 513
FROZEN_GRAPH_NAME = 'frozen_inference_graph'

def __init__(self, tarball_path):
    """    def __init__(self, tarball_path):
    """
    Creates deeplab model.
    """
    self.graph = tf.Graph()

    #graph_def = None

    graph_def = tf.GraphDef.FromString(open(tarball_path, 'rb').read())

    if graph_def is None:
        raise RuntimeError('Cannot find inference graph in tar archive.')

    with self.graph.as_default():
        tf.import_graph_def(graph_def, name='')

    self.sess = tf.Session(graph=self.graph)


def run(self, image):
    """
    Runs inference on a single image.

    Args:
    image: A PIL.Image object, raw input image.

    Returns:
    resized_image: RGB image resized from original input image.
    seg_map: Segmentation map of `resized_image`.
    """
    width, height = image.size
    resize_ratio = 1.0 * self.INPUT_SIZE / max(width, height)
    target_size = (int(resize_ratio * width), int(resize_ratio * height))
    resized_image = image.convert('RGB').resize(target_size, Image.ANTIALIAS)
    batch_seg_map = self.sess.run(self.OUTPUT_TENSOR_NAME,
                                  feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]})
    seg_map = batch_seg_map[0]
    return resized_image, seg_map

def create_pascal_label_colormap():
"""
Creates a label colormap used in PASCAL VOC segmentation benchmark.

Returns:
    A Colormap for visualizing segmentation results.
"""
colormap = np.zeros((256, 3), dtype=int)
ind = np.arange(256, dtype=int)

for shift in reversed(range(8)):
    for channel in range(3):
        colormap[:, channel] |= ((ind >> channel) & 1) << shift
    ind >>= 3

return colormap

def label_to_color_image(label):
"""
Adds color defined by the dataset colormap to the label.

Args:
    label: A 2D array with integer type, storing the segmentation label.

Returns:
    result: A 2D array with floating type. The element of the array
    is the color indexed by the corresponding element in the input label
    to the PASCAL color map.

Raises:
    ValueError: If label is not of rank 2 or its value is larger than color
    map maximum entry.
"""
if label.ndim != 2:
    raise ValueError('Expect 2-D input label')

colormap = create_pascal_label_colormap()

if np.max(label) >= len(colormap):
    raise ValueError('label value too large.')

return colormap[label]

def write_result(imagefile, imgname, outputDir):

orignal_im = Image.open(imagefile)
print('running deeplab on image %s...' % imagefile)
resized_im, seg_map = MODEL.run(orignal_im)
seg_image = label_to_color_image(seg_map).astype(np.uint8)
img_path = os.path.join(outputDir, imgname)
cv2.imwrite(img_path, seg_image)

`