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/image_prep.m b/example_images/image_prep.m new file mode 100644 index 0000000..08285c9 --- /dev/null +++ b/example_images/image_prep.m @@ -0,0 +1,21 @@ +img = imread('/home/tatsch/PSPNet-Keras-tensorflow/example_images/cityscapes.png'); +crop_size = 713; +resized = imresize(img, [crop_size crop_size], 'bilinear'); +mean_r = 123.68; %means to be subtracted and the given values are used in our training stage +mean_g = 116.779; +mean_b = 103.939; + +mean_rgb = [mean_r; mean_g; mean_b]; +mean_rgb = reshape(mean_rgb, 1, 1, 3); + +im_mean = repmat(mean_rgb, [crop_size, crop_size, 1]); + +resized_f = single(resized); +demeaned = resized_f - im_mean; + +mean_rc = mean(mean(demeaned(:, :, 1))); +mean_bc = mean(mean(demeaned(:, :, 2))); +mean_gc = mean(mean(demeaned(:, :, 3))); + +figure; +imshow(demeaned) \ No newline at end of file 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..3af91fc 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..fe58af5 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..59ee2cd 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..f88ee91 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..fb47c8a 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..2dfa339 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..fea81c8 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..932c58b 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..7f59547 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..630a075 100755 --- a/pspnet.py +++ b/pspnet.py @@ -1,7 +1,9 @@ #!/usr/bin/env python 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 +12,22 @@ from keras.models import model_from_json import tensorflow as tf import layers_builder as layers import utils +import matplotlib.pyplot as plt + # These are the means for the ImageNet pretrained ResNet DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB order 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()) @@ -42,32 +47,32 @@ class PSPNet(object): img: must be rowsxcolsx3 """ h_ori, w_ori = img.shape[:2] - - # Preprocess img = misc.imresize(img, self.input_shape) - - img = img - DATA_MEAN - img = img[:, :, ::-1] # RGB => BGR - img = img.astype('float32') - print("Predicting...") - + img = self.preprocess_image(img) 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 preprocess_image(self, img): + """Preprocess an image as input.""" + float_img = img.astype('float16') + centered_image = float_img - DATA_MEAN + return centered_image + def feed_forward(self, data): + """Pass an image through the network.""" assert data.shape == (self.input_shape[0], self.input_shape[1], 3) - data = data[np.newaxis, :, :, :] + data = data[:, :, ::-1] # RGB => BGR + data = data[np.newaxis, :, :, :] # Append sample dimension for keras # utils.debug(self.model, data) pred = self.model.predict(data) return pred[0] 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 +80,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 +93,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 +125,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 +134,58 @@ 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 sliding_prediction(full_image, net): + """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 = int(x2 - tile_size[1]) + y1 = int(y2 - tile_size[0]) + if x1 < 0: # for portrait the x1 underflows sometimes + x1 = 0 + 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) + 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 + return full_probs + + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument('-m', '--model', type=str, default='pspnet50_ade20k', @@ -132,9 +198,11 @@ 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_prediction', type=bool, default=False, + help="Whether the network should be slided along the original image for prediction.") 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 +225,13 @@ if __name__ == "__main__": else: print("Network architecture not implemented.") - probs = pspnet.predict(img) + # TODO: implement score flips + + if args.sliding_prediction: + probs = sliding_prediction(img, pspnet) + else: + probs = pspnet.predict(img) + 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[...]