diff --git a/README.md b/README.md index 3b85838..feab14e 100644 --- a/README.md +++ b/README.md @@ -1,30 +1,34 @@ # Keras implementation of [PSPNet(caffe)](https://github.com/hszhao/PSPNet) -Implemented Architecture of pyramid scene parsing network in Keras +Implemented Architecture of Pyramid Scene Parsing Network in Keras. -Converted trained weights needed to run the network. +Converted trained weights are needed to run the network. +Weights of the original caffemodel can be converted with weight_converter.py as follows: -Download converted weights here: -[link:pspnet50_ade20k.npy](https://www.dropbox.com/s/ms8afun494dlh1t/pspnet50_ade20k.npy?dl=0) - -And place in directory with pspnet50_ade20k.npy - -Weights from caffemodel were converted by, weight_converter.py. The usage of this file is ```bash python weight_converter.py ``` -Running this need to compile the original PSPNet caffe code and pycaffe. -Interpolation layer is implemented in code as custom layer "Interp" +Running this needs the compiled original PSPNet caffe code and pycaffe. +Already converted weights can be downloaded here: + +[pspnet50_ade20k.npy](https://www.dropbox.com/s/ms8afun494dlh1t/pspnet50_ade20k.npy?dl=0) +[pspnet101_cityscapes.npy](https://www.dropbox.com/s/b21j6hi6qql90l0/pspnet101_cityscapes.npy?dl=0) +[pspnet101_voc2012.npy](https://www.dropbox.com/s/xkjmghsbn6sfj9k/pspnet101_voc2012.npy?dl=0) + +weights should be placed in the directory with pspnet.py + +The interpolation layer is implemented as custom layer "Interp" ## Important Results Keras: ![Original](test.jpg) -![New](out.jpg) -![New](probs.jpg) +![New](test_seg.jpg) +![New](test_seg_blended.jpg) +![New](test_probs.jpg) ## Pycaffe result ![Pycaffe results](test_pycaffe.jpg) @@ -32,12 +36,11 @@ Results Keras: 1. Tensorflow 2. Keras 3. numpy -4. pycaffe(PSPNet)(optional) +4. pycaffe(PSPNet)(optional for converting the weights) -## Usage: +## Usage: ```bash python pspnet.py --input-path INPUT_PATH --output-path OUTPUT_PATH ``` - diff --git a/layers_builder.py b/layers_builder.py index e007117..ca7f101 100644 --- a/layers_builder.py +++ b/layers_builder.py @@ -19,7 +19,8 @@ def BN(name=""): def Interp(x, shape): from keras.backend import tf as ktf new_height, new_width = shape - resized = ktf.image.resize_images(x, [new_height, new_width], align_corners=True) + resized = ktf.image.resize_images(x, [new_height, new_width], + align_corners=True) return resized @@ -33,19 +34,23 @@ def residual_conv(prev, level, pad=1, lvl=1, sub_lvl=1, modify_stride=False): "conv"+lvl+"_" + sub_lvl + "_1x1_increase", "conv"+lvl+"_" + sub_lvl + "_1x1_increase_bn"] if modify_stride is False: - prev = Conv2D(64 * level, (1, 1), strides=(1, 1), name=names[0], use_bias=False)(prev) + prev = Conv2D(64 * level, (1, 1), strides=(1, 1), name=names[0], + use_bias=False)(prev) elif modify_stride is True: - prev = Conv2D(64 * level, (1, 1), strides=(2, 2), name=names[0], use_bias=False)(prev) + prev = Conv2D(64 * level, (1, 1), strides=(2, 2), name=names[0], + use_bias=False)(prev) prev = BN(name=names[1])(prev) prev = Activation('relu')(prev) prev = ZeroPadding2D(padding=(pad, pad))(prev) - prev = Conv2D(64 * level, (3, 3), strides=(1, 1), dilation_rate=pad, name=names[2], use_bias=False)(prev) + prev = Conv2D(64 * level, (3, 3), strides=(1, 1), dilation_rate=pad, + name=names[2], use_bias=False)(prev) prev = BN(name=names[3])(prev) prev = Activation('relu')(prev) - prev = Conv2D(256 * level, (1, 1), strides=(1, 1), name=names[4], use_bias=False)(prev) + prev = Conv2D(256 * level, (1, 1), strides=(1, 1), name=names[4], + use_bias=False)(prev) prev = BN(name=names[5])(prev) return prev @@ -57,9 +62,11 @@ def short_convolution_branch(prev, level, lvl=1, sub_lvl=1, modify_stride=False) "conv" + lvl+"_" + sub_lvl + "_1x1_proj_bn"] if modify_stride is False: - prev = Conv2D(256 * level, (1, 1), strides=(1, 1), name=names[0], use_bias=False)(prev) + prev = Conv2D(256 * level, (1, 1), strides=(1, 1), name=names[0], + use_bias=False)(prev) elif modify_stride is True: - prev = Conv2D(256 * level, (1, 1), strides=(2, 2), name=names[0], use_bias=False)(prev) + prev = Conv2D(256 * level, (1, 1), strides=(2, 2), name=names[0], + use_bias=False)(prev) prev = BN(name=names[1])(prev) return prev @@ -85,7 +92,8 @@ def residual_short(prev_layer, level, pad=1, lvl=1, sub_lvl=1, modify_stride=Fal def residual_empty(prev_layer, level, pad=1, lvl=1, sub_lvl=1): prev_layer = Activation('relu')(prev_layer) - block_1 = residual_conv(prev_layer, level, pad=pad, lvl=lvl, sub_lvl=sub_lvl) + block_1 = residual_conv(prev_layer, level, pad=pad, + lvl=lvl, sub_lvl=sub_lvl) block_2 = empty_branch(prev_layer) added = Add()([block_1, block_2]) return added @@ -102,19 +110,23 @@ def ResNet(inp, layers): # Short branch(only start of network) - cnv1 = Conv2D(64, (3, 3), strides=(2, 2), padding='same', name=names[0], use_bias=False)(inp) # "conv1_1_3x3_s2" + cnv1 = Conv2D(64, (3, 3), strides=(2, 2), padding='same', name=names[0], + use_bias=False)(inp) # "conv1_1_3x3_s2" bn1 = BN(name=names[1])(cnv1) # "conv1_1_3x3_s2/bn" relu1 = Activation('relu')(bn1) # "conv1_1_3x3_s2/relu" - cnv1 = Conv2D(64, (3, 3), strides=(1, 1), padding='same', name=names[2], use_bias=False)(relu1) # "conv1_2_3x3" + 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" - cnv1 = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name=names[4], use_bias=False)(relu1) # "conv1_3_3x3" + cnv1 = Conv2D(128, (3, 3), strides=(1, 1), padding='same', name=names[4], + use_bias=False)(relu1) # "conv1_3_3x3" bn1 = BN(name=names[5])(cnv1) # "conv1_3_3x3/bn" relu1 = Activation('relu')(bn1) # "conv1_3_3x3/relu" - res = MaxPooling2D(pool_size=(3, 3), padding='same', strides=(2, 2))(relu1) # "pool1_3x3_s2" + res = MaxPooling2D(pool_size=(3, 3), padding='same', + strides=(2, 2))(relu1) # "pool1_3x3_s2" # ---Residual layers(body of network) @@ -166,7 +178,8 @@ def interp_block(prev_layer, level, feature_map_shape, str_lvl=1, ): kernel = (10*level, 10*level) strides = (10*level, 10*level) prev_layer = AveragePooling2D(kernel, strides=strides)(prev_layer) - prev_layer = Conv2D(512, (1, 1), strides=(1, 1), name=names[0], use_bias=False)(prev_layer) + prev_layer = Conv2D(512, (1, 1), strides=(1, 1), name=names[0], + use_bias=False)(prev_layer) prev_layer = BN(name=names[1])(prev_layer) prev_layer = Activation('relu')(prev_layer) prev_layer = Lambda(Interp, arguments={'shape': feature_map_shape})(prev_layer) @@ -201,7 +214,8 @@ def build_pspnet(nb_classes, resnet_layers, input_shape, activation='softmax'): res = ResNet(inp, layers=resnet_layers) psp = PSPNet(res, input_shape) - x = Conv2D(512, (3, 3), strides=(1, 1), padding="same", name="conv5_4", use_bias=False)(psp) + x = Conv2D(512, (3, 3), strides=(1, 1), padding="same", name="conv5_4", + use_bias=False)(psp) x = BN(name="conv5_4_bn")(x) x = Activation('relu')(x) x = Dropout(0.1)(x) diff --git a/out.jpg b/out.jpg deleted file mode 100644 index 14c2689..0000000 Binary files a/out.jpg and /dev/null differ diff --git a/probs.jpg b/probs.jpg deleted file mode 100644 index 9bc9563..0000000 Binary files a/probs.jpg and /dev/null differ diff --git a/pspnet.py b/pspnet.py index 1bfa6ed..2635d91 100644 --- a/pspnet.py +++ b/pspnet.py @@ -10,8 +10,8 @@ import tensorflow as tf import layers_builder as layers import utils - -DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB, these are the means for the ImageNet pretrained ResNet +# These are the means for the ImageNet pretrained ResNet +DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB order class PSPNet(object): @@ -27,8 +27,10 @@ class PSPNet(object): self.model = model_from_json(file_handle.read()) self.model.load_weights(h5_path) else: - print("No Keras model & weights found, importing from numpy weights.") - self.model = layers.build_pspnet(nb_classes=nb_classes, resnet_layers=resnet_layers, input_shape=self.input_shape) + print("No Keras model & weights found, import from npy weights.") + self.model = layers.build_pspnet(nb_classes=nb_classes, + resnet_layers=resnet_layers, + input_shape=self.input_shape) self.set_npy_weights(weights) def predict(self, img): @@ -50,7 +52,8 @@ class PSPNet(object): 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) + probs = ndimage.zoom(probs, (1.*h_ori/h, 1.*w_ori/w, 1.), + order=1, prefilter=False) print("Finished prediction...") return probs @@ -79,7 +82,8 @@ class PSPNet(object): scale = weights[layer.name]['scale'].reshape(-1) offset = weights[layer.name]['offset'].reshape(-1) - self.model.get_layer(layer.name).set_weights([mean, variance, scale, offset]) + self.model.get_layer(layer.name).set_weights([mean, variance, + scale, offset]) elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu': try: @@ -87,7 +91,8 @@ class PSPNet(object): self.model.get_layer(layer.name).set_weights([weight]) except Exception as err: biases = weights[layer.name]['biases'] - self.model.get_layer(layer.name).set_weights([weight, biases]) + self.model.get_layer(layer.name).set_weights([weight, + biases]) print('Finished importing weights.') print("Writing keras model & weights") @@ -102,21 +107,29 @@ class PSPNet50(PSPNet): """Build a PSPNet based on a 50-Layer ResNet.""" def __init__(self, nb_classes, weights, input_shape): - PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=50, input_shape=input_shape, weights=weights) + PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=50, + input_shape=input_shape, weights=weights) class PSPNet101(PSPNet): """Build a PSPNet based on a 101-Layer ResNet.""" def __init__(self, nb_classes, weights, input_shape): - PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=101, input_shape=input_shape, weights=weights) + PSPNet.__init__(self, nb_classes=nb_classes, resnet_layers=101, + input_shape=input_shape, weights=weights) if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('-m', '--model', type=str, default='pspnet50_ade20k', help='Model/Weights to use', choices=['pspnet50_ade20k', 'pspnet101_cityscapes', 'pspnet101_voc2012']) - parser.add_argument('-i', '--input_path', type=str, default='test.jpg', help='Path the input image') - parser.add_argument('-o', '--output_path', type=str, default='test.jpg', help='Path to output') + parser.add_argument('-m', '--model', type=str, default='pspnet50_ade20k', + help='Model/Weights to use', + choices=['pspnet50_ade20k', + 'pspnet101_cityscapes', + 'pspnet101_voc2012']) + parser.add_argument('-i', '--input_path', type=str, default='test.jpg', + help='Path the input image') + parser.add_argument('-o', '--output_path', type=str, default='test.jpg', + help='Path to output') parser.add_argument('--id', default="0") args = parser.parse_args() @@ -130,12 +143,15 @@ if __name__ == "__main__": print(args) if "pspnet50" in args.model: - pspnet = PSPNet50(nb_classes=150, input_shape=(473, 473), weights=args.model) + pspnet = PSPNet50(nb_classes=150, input_shape=(473, 473), + weights=args.model) elif "pspnet101" in args.model: if "cityscapes" in args.model: - pspnet = PSPNet101(nb_classes=19, input_shape=(713, 713), weights=args.model) + pspnet = PSPNet101(nb_classes=19, input_shape=(713, 713), + weights=args.model) if "voc2012" in args.model: - pspnet = PSPNet101(nb_classes=21, input_shape=(473, 473), weights=args.model) + pspnet = PSPNet101(nb_classes=21, input_shape=(473, 473), + weights=args.model) else: print("Network architecture not implemented.") @@ -146,7 +162,8 @@ if __name__ == "__main__": cm = np.argmax(probs, axis=2) + 1 pm = np.max(probs, axis=2) color_cm = utils.add_color(cm) - alpha_blended = 0.5 * color_cm * 255 + 0.5 * img # color cm is [0.0-1.0] img [0-255] + # color cm is [0.0-1.0] img is [0-255] + alpha_blended = 0.5 * color_cm * 255 + 0.5 * img filename, ext = splitext(args.output_path) misc.imsave(filename + "_seg" + ext, color_cm) misc.imsave(filename + "_probs" + ext, pm) diff --git a/utils.py b/utils.py index 8f179f5..75e3dca 100644 --- a/utils.py +++ b/utils.py @@ -32,8 +32,9 @@ def print_activation(model, layer_name, data): intermediate_layer_model = Model(inputs=model.input, outputs=model.get_layer(layer_name).output) io = intermediate_layer_model.predict(data) - print (layer_name, array_to_str(io)) + print(layer_name, array_to_str(io)) def array_to_str(a): - return "{} {} {} {} {}".format(a.dtype, a.shape, np.min(a), np.max(a), np.mean(a)) + return "{} {} {} {} {}".format(a.dtype, a.shape, np.min(a), + np.max(a), np.mean(a))