Deeplab - Inconsistent Inference vs Visualization Performance on trained Deeplab model

Question

Describe the problem

I have successfully trained my model on a custom dataset with 4 classes of size 480x640, with an xception65 encoder, using Deeplab. I am getting decent results on the validation set whenever I use the vis.py script: EvalImageA_ckpt, EvalImageB_ckpt. However, I am not getting the same results on the same images when I freeze the model.

I froze the model using export_model.py and successfully outputted a frozen_model.pb file. However, when I run inferences using this pb file, the outputs are always 0 (i.e. everything is classified as "background") on the same exact images I provided links to above. Everything is black!

I believe this to be an issue with how I am exporting or loading the model, and not necessarily with the model itself because the performance on the images is different between running the vis.py script and my custom code for inference. Perhaps I am not loading the graph or initializing the variables correctly. Or perhaps I'm not saving the weights correctly in the first place. Any help would be greatly appreciated!

Source code

Below I provide my code for inference:

from deeplab.utils import get_dataset_colormap
from PIL import Image
import tensorflow as tf
import time
import matplotlib.pyplot as plt
import numpy as np
import cv2
import os
import glob


# tensorflow arguments
flags = tf.app.flags  # flag object for setup
FLAGS = flags.FLAGS   # object to access initialized flags
flags.DEFINE_string('frozen', None,
                    'The path/to/frozen.pb file.')

def _load_graph(frozen):
    print('Loading model `deeplabv3_graph` into memory from',frozen)
    with tf.gfile.GFile(frozen, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(
            graph_def, 
            input_map=None, 
            return_elements=None, 
            name="", 
            op_dict=None, 
            producer_op_list=None
        )
    return graph

def _run_inferences(sess, image, title):
    batch_seg_map = sess.run('SemanticPredictions:0',
        feed_dict={'ImageTensor:0': [np.asarray(image)]})
    semantic_prediction = get_dataset_colormap.label_to_color_image(batch_seg_map[0],
        dataset=get_dataset_colormap.__PRDL3_V1).astype(np.uint8)
    plt.imshow(semantic_prediction)
    plt.axis('off')
    plt.title(title)
    plt.show()


def main(argv):
    # initialize model
    frozen = os.path.normpath(FLAGS.frozen)
    assert os.path.isfile(frozen)
    graph = _load_graph(frozen)

    # open graph resource and begin inference in-loop
    with tf.Session(graph=graph) as sess:
        for img_path in glob.glob('*.png'):
            img = Image.open(img_path).convert('RGB')
            _run_inferences(sess, img, img_path)

if __name__ == '__main__':
    flags.mark_flag_as_required('frozen')
    tf.app.run()  # call the main() function

And below is my code for exporting the model, using the provided export_model.py script.

python export_model.py \
--logtostderr \
--atrous_rates=6 \
--atrous_rates=12 \
--atrous_rates=18 \
--output_stride=16 \
--checkpoint_path="/path/to/.../model.ckpt-32245" \
--export_path="/path/to/.../frozen_4_11_19.pb" \
--model_variant="xception_65" \
--num_classes=4 \
--crop_size=481 \
--crop_size=641 \
--inference_scales=1.0

System information

• What is the top-level directory of the model you are using: deeplab
• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 10 Enterprise
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 1.12.0
• Bazel version (if compiling from source): N/A
• CUDA/cuDNN version: 9
• GPU model and memory: NVIDIA Quadro M4000, 8GB
• Exact command to reproduce: Does not apply

Answer 1

Run export model script with following flags

python deeplab/export_model.py \
--logtostderr \
--model_variant="xception_65" \
--atrous_rates=6 \
--atrous_rates=12 \
--atrous_rates=18 \
--output_stride=16 \
--num_classes=2 \
--decoder_output_stride=4 \
--crop_size=<<crop width>> \
--crop_size=<<crop height>> \
--dataset="shoes" \
--checkpoint_path="<<Checkpoint path>>" \
--export_path="<<Output frozen graph path>>" \

Answer 2

I'm also struggling with my inferences results. But in my case I have quite satisfing results when using the exported model, it's just that they are less accurate than my visualization results.

Here is my script, it's a script based in the one available as demo. Hope it helps

import os

from matplotlib import gridspec
from matplotlib import pyplot as plt
import numpy as np
from PIL import Image
import time
import cv2
from tqdm import tqdm

import tensorflow as tf

# Needed to show segmentation colormap labels

from deeplab.utils import get_dataset_colormap
from deeplab.utils import labels_alstom

flags = tf.app.flags

FLAGS = flags.FLAGS

flags.DEFINE_string('model_dir', None, 'Where the model is')
flags.DEFINE_string('image_dir', None, 'Where the image is')
flags.DEFINE_string('save_dir', None, 'Dir for saving results')
flags.DEFINE_string('image_name', None, 'Image name')



class DeepLabModel(object):
    """Class to load deeplab model and run inference."""

    INPUT_TENSOR_NAME = 'ImageTensor:0'
    OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
    INPUT_SIZE = 513

    def __init__(self, model_dir):
        """Creates and loads pretrained deeplab model."""
        self.graph = tf.Graph()

        graph_def = None
        # Extract frozen graph from tar archive.
        model_filename = FLAGS.model_dir
        with tf.gfile.FastGFile(model_filename, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.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)
        print('Image resized')
        start_time = time.time()
        batch_seg_map = self.sess.run(
            self.OUTPUT_TENSOR_NAME,
            feed_dict={self.INPUT_TENSOR_NAME: [np.asarray(resized_image)]})
        print('Image processing finished')
        print('Elapsed time : ' + str(time.time() - start_time))
        seg_map = batch_seg_map[0]
        return resized_image, seg_map


model = DeepLabModel(FLAGS.model_dir)
print('Model created successfully')



def vis_segmentation(image, seg_map):
    
    seg_image = get_dataset_colormap.label_to_color_image(
         seg_map, get_dataset_colormap.get_alstom_name()).astype(np.uint8)
         
    return seg_image



def run_demo_image(image_path):
    try:
        print(image_path)
        orignal_im = Image.open(image_path)

    except IOError:
        print ('Failed to read image from %s.' % image_path)
        return
    print ('running deeplab on image...')
    resized_im, seg_map = model.run(orignal_im)

    return vis_segmentation(resized_im, seg_map)



IMAGE_DIR = FLAGS.image_dir


files = os.listdir(FLAGS.image_dir)
for f in tqdm(files):

    prediction = run_demo_image(IMAGE_DIR+f)
    Image.fromarray(prediction).save(FLAGS.save_dir+'prediction_'+f)