Merge pull request #13 from jmtatsch/master

Added sliding evaluation for pspnet

README.md

@@ -21,9 +21,7 @@ npy weights should be placed in the directory weights/npy.
 
 The interpolation layer is implemented as custom layer "Interp"
 
-## Important
-
-Results Keras:
+## Keras result:
 ![Original](example_images/ade20k.jpg)
 ![New](example_results/ade20k_seg.jpg)
 ![New](example_results/ade20k_seg_blended.jpg)
@@ -39,14 +37,21 @@ Results Keras:
 ![New](example_results/pascal_voc_seg_blended.jpg)
 ![New](example_results/pascal_voc_probs.jpg)
 
-## Pycaffe result
+## Pycaffe result:
 ![Pycaffe results](example_results/ade20k_seg_pycaffe.jpg)
+
 ## Dependencies:
-1. Tensorflow
+
+
+1. Tensorflow (-gpu)
 2. Keras
 3. numpy
+4. scipy
 4. pycaffe(PSPNet)(optional for converting the weights)
 
+```bash
+pip install -r requirements.txt --upgrade
+```
 
 ## Usage:
 

layers_builder.py

@@ -6,8 +6,6 @@ from keras.layers.merge import Concatenate, Add
 from keras.models import Model
 from keras.optimizers import SGD
 
-import tensorflow as tf
-
 learning_rate = 1e-3  # Layer specific learning rate
 # Weight decay not implemented
 
@@ -118,7 +116,7 @@ def ResNet(inp, layers):
     cnv1 = Conv2D(64, (3, 3), strides=(1, 1), padding='same', name=names[2],
                   use_bias=False)(relu1)  # "conv1_2_3x3"
     bn1 = BN(name=names[3])(cnv1)  # "conv1_2_3x3/bn"
-    relu1 = Activation('relu')(bn1)                 # "conv1_2_3x3/relu"
+    relu1 = Activation('relu')(bn1)  # "conv1_2_3x3/relu"
 
     cnv1 = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name=names[4],
                   use_bias=False)(relu1)  # "conv1_3_3x3"
@@ -167,7 +165,6 @@ def ResNet(inp, layers):
 
 
 def interp_block(prev_layer, level, feature_map_shape, str_lvl=1, ):
-
     str_lvl = str(str_lvl)
 
     names = [
@@ -186,7 +183,7 @@ def interp_block(prev_layer, level, feature_map_shape, str_lvl=1, ):
     return prev_layer
 
 
-def PSPNet(res, input_shape):
+def build_pyramid_pooling_module(res, input_shape):
     """Build the Pyramid Pooling Module."""
     # ---PSPNet concat layers with Interpolation
     feature_map_size = tuple(int(ceil(input_dim / 8.0)) for input_dim in input_shape)
@@ -197,7 +194,7 @@ def PSPNet(res, input_shape):
     interp_block3 = interp_block(res, 2, feature_map_size, str_lvl=3)
     interp_block6 = interp_block(res, 1, feature_map_size, str_lvl=6)
 
-    # concat all these layers. resulted shape=(1,60,60,4096)
+    # concat all these layers. resulted shape=(1,feature_map_size_x,feature_map_size_y,4096)
     res = Concatenate()([res,
                          interp_block6,
                          interp_block3,
@@ -212,7 +209,7 @@ def build_pspnet(nb_classes, resnet_layers, input_shape, activation='softmax'):
 
     inp = Input((input_shape[0], input_shape[1], 3))
     res = ResNet(inp, layers=resnet_layers)
-    psp = PSPNet(res, input_shape)
+    psp = build_pyramid_pooling_module(res, input_shape)
 
     x = Conv2D(512, (3, 3), strides=(1, 1), padding="same", name="conv5_4",
                use_bias=False)(psp)

pspnet.py

@@ -1,7 +1,14 @@
 #!/usr/bin/env python
+"""
+This module is a Keras/Tensorflow based implementation of Pyramid Scene Parsing Networks.
+
+Original paper & code published by Hengshuang Zhao et al. (2017)
+"""
 from __future__ import print_function
-import os
-from os.path import splitext, join
+from __future__ import division
+from os.path import splitext, join, isfile
+from os import environ
+from math import ceil
 import argparse
 import numpy as np
 from scipy import misc, ndimage
@@ -10,19 +17,25 @@ from keras.models import model_from_json
 import tensorflow as tf
 import layers_builder as layers
 import utils
+import matplotlib.pyplot as plt
+
+__author__ = "Vlad Kryvoruchko, Chaoyue Wang, Jeffrey Hu & Julian Tatsch"
+
 
 # These are the means for the ImageNet pretrained ResNet
 DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]])  # RGB order
+EVALUATION_SCALES = [1.0]  # must be all floats!
 
 
 class PSPNet(object):
-    """Pyramid Scene Parsing Network by Hengshuang Zhao et al 2017"""
+    """Pyramid Scene Parsing Network by Hengshuang Zhao et al 2017."""
 
     def __init__(self, nb_classes, resnet_layers, input_shape, weights):
+        """Instanciate a PSPNet."""
         self.input_shape = input_shape
         json_path = join("weights", "keras", weights + ".json")
         h5_path = join("weights", "keras", weights + ".h5")
-        if os.path.isfile(json_path) and os.path.isfile(h5_path):
+        if isfile(json_path) and isfile(h5_path):
             print("Keras model & weights found, loading...")
             with open(json_path, 'r') as file_handle:
                 self.model = model_from_json(file_handle.read())
@@ -34,7 +47,7 @@ class PSPNet(object):
                                              input_shape=self.input_shape)
             self.set_npy_weights(weights)
 
-    def predict(self, img):
+    def predict(self, img, flip_evaluation):
         """
         Predict segementation for an image.
 
@@ -42,32 +55,36 @@ class PSPNet(object):
             img: must be rowsxcolsx3
         """
         h_ori, w_ori = img.shape[:2]
+        if img.shape[0:2] != self.input_shape:
+            print("Input %s not fitting for network size %s, resizing. You may want to try sliding prediction for better results." % (img.shape[0:2], self.input_shape))
+            img = misc.imresize(img, self.input_shape)
+        input_data = self.preprocess_image(img)
+        # utils.debug(self.model, input_data)
 
-        # Preprocess
-        img = misc.imresize(img, self.input_shape)
+        regular_prediction = self.model.predict(input_data)[0]
+        if flip_evaluation:
+            print("Predict flipped")
+            flipped_prediction = np.fliplr(self.model.predict(np.flip(input_data, axis=2))[0])
+            prediction = (regular_prediction + flipped_prediction)
+        else:
+            prediction = regular_prediction
 
-        img = img - DATA_MEAN
-        img = img[:, :, ::-1]  # RGB => BGR
-        img = img.astype('float32')
-        print("Predicting...")
+        if img.shape[0:1] != self.input_shape:  # upscale prediction if necessary
+            h, w = prediction.shape[:2]
+            prediction = ndimage.zoom(prediction, (1.*h_ori/h, 1.*w_ori/w, 1.),
+                                      order=1, prefilter=False)
+        return prediction
 
-        probs = self.feed_forward(img)
-        h, w = probs.shape[:2]
-        probs = ndimage.zoom(probs, (1.*h_ori/h, 1.*w_ori/w, 1.),
-                             order=1, prefilter=False)
-        print("Finished prediction...")
-
-        return probs
-
-    def feed_forward(self, data):
-        assert data.shape == (self.input_shape[0], self.input_shape[1], 3)
-        data = data[np.newaxis, :, :, :]
-
-        # utils.debug(self.model, data)
-        pred = self.model.predict(data)
-        return pred[0]
+    def preprocess_image(self, img):
+        """Preprocess an image as input."""
+        float_img = img.astype('float16')
+        centered_image = float_img - DATA_MEAN
+        bgr_image = centered_image[:, :, ::-1]  # RGB => BGR
+        input_data = bgr_image[np.newaxis, :, :, :]  # Append sample dimension for keras
+        return input_data
 
     def set_npy_weights(self, weights_path):
+        """Set weights from the intermediary npy file."""
         npy_weights_path = join("weights", "npy", weights_path + ".npy")
         json_path = join("weights", "keras", weights_path + ".json")
         h5_path = join("weights", "keras", weights_path + ".h5")
@@ -75,8 +92,11 @@ class PSPNet(object):
         print("Importing weights from %s" % npy_weights_path)
         weights = np.load(npy_weights_path).item()
 
+        whitelist = ["InputLayer", "Activation", "ZeroPadding2D", "Add", "MaxPooling2D", "AveragePooling2D", "Lambda", "Concatenate", "Dropout"]
+
+        weights_set = 0
         for layer in self.model.layers:
-            print(layer.name)
+            print("Processing %s" % layer.name)
             if layer.name[:4] == 'conv' and layer.name[-2:] == 'bn':
                 mean = weights[layer.name]['mean'].reshape(-1)
                 variance = weights[layer.name]['variance'].reshape(-1)
@@ -85,15 +105,24 @@ class PSPNet(object):
 
                 self.model.get_layer(layer.name).set_weights([mean, variance,
                                                              scale, offset])
-
+                weights_set += 1
             elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu':
                 try:
                     weight = weights[layer.name]['weights']
                     self.model.get_layer(layer.name).set_weights([weight])
-                except Exception as err:
+                except Exception:
                     biases = weights[layer.name]['biases']
                     self.model.get_layer(layer.name).set_weights([weight,
                                                                  biases])
+                weights_set += 1
+            elif layer.__class__.__name__ in whitelist:
+                # print("Nothing to set in %s" % layer.__class__.__name__)
+                pass
+            else:
+                print("Warning: Did not find weights for keras layer %s in numpy weights" % layer)
+
+        print("Set a total of %i weights" % weights_set)
+
         print('Finished importing weights.')
 
         print("Writing keras model & weights")
@@ -108,6 +137,7 @@ class PSPNet50(PSPNet):
     """Build a PSPNet based on a 50-Layer ResNet."""
 
     def __init__(self, nb_classes, weights, input_shape):
+        """Instanciate a PSPNet50."""
         PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=50,
                         input_shape=input_shape, weights=weights)
 
@@ -116,10 +146,91 @@ class PSPNet101(PSPNet):
     """Build a PSPNet based on a 101-Layer ResNet."""
 
     def __init__(self, nb_classes, weights, input_shape):
+        """Instanciate a PSPNet101."""
         PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=101,
                         input_shape=input_shape, weights=weights)
 
 
+def pad_image(img, target_size):
+    """Pad an image up to the target size."""
+    rows_missing = target_size[0] - img.shape[0]
+    cols_missing = target_size[1] - img.shape[1]
+    padded_img = np.pad(img, ((0, rows_missing), (0, cols_missing), (0, 0)), 'constant')
+    return padded_img
+
+
+def visualize_prediction(prediction):
+    """Visualize prediction."""
+    cm = np.argmax(prediction, axis=2) + 1
+    color_cm = utils.add_color(cm)
+    plt.imshow(color_cm)
+    plt.show()
+
+
+def predict_sliding(full_image, net, flip_evaluation):
+    """Predict on tiles of exactly the network input shape so nothing gets squeezed."""
+    tile_size = net.input_shape
+    classes = net.model.outputs[0].shape[3]
+    overlap = 1/3
+
+    stride = ceil(tile_size[0] * (1 - overlap))
+    tile_rows = int(ceil((full_image.shape[0] - tile_size[0]) / stride) + 1)  # strided convolution formula
+    tile_cols = int(ceil((full_image.shape[1] - tile_size[1]) / stride) + 1)
+    print("Need %i x %i prediction tiles @ stride %i px" % (tile_cols, tile_rows, stride))
+    full_probs = np.zeros((full_image.shape[0], full_image.shape[1], classes))
+    count_predictions = np.zeros((full_image.shape[0], full_image.shape[1], classes))
+    tile_counter = 0
+    for row in range(tile_rows):
+        for col in range(tile_cols):
+            x1 = int(col * stride)
+            y1 = int(row * stride)
+            x2 = min(x1 + tile_size[1], full_image.shape[1])
+            y2 = min(y1 + tile_size[0], full_image.shape[0])
+            x1 = max(int(x2 - tile_size[1]), 0)  # for portrait images the x1 underflows sometimes
+            y1 = max(int(y2 - tile_size[0]), 0)  # for very few rows y1 underflows
+
+            img = full_image[y1:y2, x1:x2]
+            padded_img = pad_image(img, tile_size)
+            # plt.imshow(padded_img)
+            # plt.show()
+            tile_counter += 1
+            print("Predicting tile %i" % tile_counter)
+            padded_prediction = net.predict(padded_img, flip_evaluation)
+            prediction = padded_prediction[0:img.shape[0], 0:img.shape[1], :]
+            count_predictions[y1:y2, x1:x2] += 1
+            full_probs[y1:y2, x1:x2] += prediction  # accumulate the predictions also in the overlapping regions
+
+    # average the predictions in the overlapping regions
+    full_probs /= count_predictions
+    # visualize normalization Weights
+    # plt.imshow(np.mean(count_predictions, axis=2))
+    # plt.show()
+    return full_probs
+
+
+def predict_multi_scale(full_image, net, scales, sliding_evaluation, flip_evaluation):
+    """Predict an image by looking at it with different scales."""
+    classes = net.model.outputs[0].shape[3]
+    full_probs = np.zeros((full_image.shape[0], full_image.shape[1], classes))
+    h_ori, w_ori = full_image.shape[:2]
+    for scale in scales:
+        print("Predicting image scaled by %f" % scale)
+        scaled_img = misc.imresize(full_image, size=scale, interp="bilinear")
+        if sliding_evaluation:
+            scaled_probs = predict_sliding(scaled_img, net, flip_evaluation)
+        else:
+            scaled_probs = net.predict(scaled_img, flip_evaluation)
+        # scale probs up to full size
+        h, w = scaled_probs.shape[:2]
+        probs = ndimage.zoom(scaled_probs, (1.*h_ori/h, 1.*w_ori/w, 1.),  # FIXME: must scale up exactly to full_image.shape
+                             order=1, prefilter=False)
+        # visualize_prediction(probs)
+        # integrate probs over all scales
+        full_probs += probs
+    full_probs /= len(scales)
+    return full_probs
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument('-m', '--model', type=str, default='pspnet50_ade20k',
@@ -132,9 +243,15 @@ if __name__ == "__main__":
     parser.add_argument('-o', '--output_path', type=str, default='example_results/ade20k.jpg',
                         help='Path to output')
     parser.add_argument('--id', default="0")
+    parser.add_argument('-s', '--sliding', action='store_true',
+                        help="Whether the network should be slided over the original image for prediction.")
+    parser.add_argument('-f', '--flip', action='store_true',
+                        help="Whether the network should predict on both image and flipped image.")
+    parser.add_argument('-ms', '--multi_scale', action='store_true',
+                        help="Whether the network should predict on multiple scales.")
     args = parser.parse_args()
 
-    os.environ["CUDA_VISIBLE_DEVICES"] = args.id
+    environ["CUDA_VISIBLE_DEVICES"] = args.id
 
     sess = tf.Session()
     K.set_session(sess)
@@ -157,7 +274,11 @@ if __name__ == "__main__":
         else:
             print("Network architecture not implemented.")
 
-        probs = pspnet.predict(img)
+        if args.multi_scale:
+            EVALUATION_SCALES = [0.5, 0.75, 1.0, 1.25, 1.5, 1.75]  # must be all floats!
+
+        probs = predict_multi_scale(img, pspnet, EVALUATION_SCALES, args.sliding, args.flip)
+
         print("Writing results...")
 
         cm = np.argmax(probs, axis=2) + 1

requirements.txt

@@ -0,0 +1,5 @@
+numpy
+scipy
+tensorflow
+tensorflow-gpu
+keras

weight_converter.py

@@ -34,7 +34,7 @@ for k, v in net.params.items():
         W = np.transpose(W, (2, 3, 1, 0))
         b = v[1].data[...]
         weights[k] = {"weights": W, "biases": b}
-    elif len(v) == 4:
+    elif len(v) == 4:  # Batchnorm layer
         k = k.replace('/', '_')
         mean = v[0].data[...]
         variance = v[1].data[...]