diff --git a/README.md b/README.md index 71d7730..6eb064d 100644 --- a/README.md +++ b/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: diff --git a/example_images/test_probs.jpg b/example_images/test_probs.jpg deleted file mode 100644 index c2e0c1a..0000000 Binary files a/example_images/test_probs.jpg and /dev/null differ diff --git a/example_images/test_seg.jpg b/example_images/test_seg.jpg deleted file mode 100644 index e9a8b0e..0000000 Binary files a/example_images/test_seg.jpg and /dev/null differ diff --git a/example_images/test_seg_blended.jpg b/example_images/test_seg_blended.jpg deleted file mode 100644 index 333825f..0000000 Binary files a/example_images/test_seg_blended.jpg and /dev/null differ diff --git a/example_results/ade20k_probs.jpg b/example_results/ade20k_probs.jpg index c2e0c1a..22000ee 100644 Binary files a/example_results/ade20k_probs.jpg and b/example_results/ade20k_probs.jpg differ diff --git a/example_results/ade20k_seg.jpg b/example_results/ade20k_seg.jpg index e9a8b0e..a294183 100644 Binary files a/example_results/ade20k_seg.jpg and b/example_results/ade20k_seg.jpg differ diff --git a/example_results/ade20k_seg_blended.jpg b/example_results/ade20k_seg_blended.jpg index 333825f..5eb0d95 100644 Binary files a/example_results/ade20k_seg_blended.jpg and b/example_results/ade20k_seg_blended.jpg differ diff --git a/example_results/cityscapes_probs.jpg b/example_results/cityscapes_probs.jpg index 8c9f1b5..050893a 100644 Binary files a/example_results/cityscapes_probs.jpg and b/example_results/cityscapes_probs.jpg differ diff --git a/example_results/cityscapes_seg.jpg b/example_results/cityscapes_seg.jpg index c542b56..f572e07 100644 Binary files a/example_results/cityscapes_seg.jpg and b/example_results/cityscapes_seg.jpg differ diff --git a/example_results/cityscapes_seg_blended.jpg b/example_results/cityscapes_seg_blended.jpg index d653f7d..809fa48 100644 Binary files a/example_results/cityscapes_seg_blended.jpg and b/example_results/cityscapes_seg_blended.jpg differ diff --git a/example_results/pascal_voc_probs.jpg b/example_results/pascal_voc_probs.jpg index 3beec57..a07408e 100644 Binary files a/example_results/pascal_voc_probs.jpg and b/example_results/pascal_voc_probs.jpg differ diff --git a/example_results/pascal_voc_seg.jpg b/example_results/pascal_voc_seg.jpg index 061dfb5..27d596e 100644 Binary files a/example_results/pascal_voc_seg.jpg and b/example_results/pascal_voc_seg.jpg differ diff --git a/example_results/pascal_voc_seg_blended.jpg b/example_results/pascal_voc_seg_blended.jpg index e36c3be..5b942e0 100644 Binary files a/example_results/pascal_voc_seg_blended.jpg and b/example_results/pascal_voc_seg_blended.jpg differ diff --git a/layers_builder.py b/layers_builder.py index ca7f101..939451b 100644 --- a/layers_builder.py +++ b/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) diff --git a/pspnet.py b/pspnet.py index 6db4191..f9122c4 100755 --- a/pspnet.py +++ b/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 diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..41661c8 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +numpy +scipy +tensorflow +tensorflow-gpu +keras diff --git a/weight_converter.py b/weight_converter.py index f775c37..738baad 100644 --- a/weight_converter.py +++ b/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[...]