mirror of
https://github.com/wassname/PSPNet-Keras-tensorflow.git
synced 2026-09-11 11:50:41 +08:00
Added support for pspnet100_cityscapes, alpha blending, keras saving & loading and much more
This commit is contained in:
Julian Tatsch
@@ -1,79 +1,122 @@
|
||||
from __future__ import print_function
|
||||
import os
|
||||
from os.path import splitext
|
||||
import argparse
|
||||
import numpy as np
|
||||
from scipy import misc, ndimage
|
||||
|
||||
from keras import backend as K
|
||||
from keras.models import model_from_json
|
||||
import tensorflow as tf
|
||||
|
||||
import layers_builder as layers
|
||||
import utils
|
||||
|
||||
WEIGHTS = 'pspnet50_ade20k.npy'
|
||||
DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB
|
||||
|
||||
class PSPNet:
|
||||
DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB, these are the means for the ImageNet pretrained ResNet
|
||||
|
||||
def __init__(self):
|
||||
self.model = layers.build_pspnet()
|
||||
set_npy_weights(self.model, WEIGHTS)
|
||||
|
||||
class PSPNet(object):
|
||||
"""Pyramid Scene Parsing Network by Hengshuang Zhao et al 2017"""
|
||||
|
||||
def __init__(self, nb_classes, resnet_layers, input_shape, weights):
|
||||
self.input_shape = input_shape
|
||||
json_path = weights + ".json"
|
||||
h5_path = weights + ".h5"
|
||||
if os.path.isfile(json_path) and os.path.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())
|
||||
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)
|
||||
self.set_npy_weights(weights)
|
||||
|
||||
def predict(self, img):
|
||||
'''
|
||||
"""
|
||||
Predict segementation for an image.
|
||||
|
||||
Arguments:
|
||||
img: must be 473x473x3
|
||||
'''
|
||||
h_ori,w_ori = img.shape[:2]
|
||||
img: must be rowsxcolsx3
|
||||
"""
|
||||
h_ori, w_ori = img.shape[:2]
|
||||
|
||||
# Preprocess
|
||||
img = misc.imresize(img, (473, 473))
|
||||
img = misc.imresize(img, self.input_shape)
|
||||
|
||||
img = img - DATA_MEAN
|
||||
img = img[:,:,::-1] # RGB => BGR
|
||||
img = img[:, :, ::-1] # RGB => BGR
|
||||
img = img.astype('float32')
|
||||
print("Predicting...")
|
||||
|
||||
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)
|
||||
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 predict_sliding_window(self, img):
|
||||
pass
|
||||
|
||||
def feed_forward(self, data):
|
||||
assert data.shape == (473,473,3)
|
||||
data = data[np.newaxis,:,:,:]
|
||||
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 set_npy_weights(model, npy_weights):
|
||||
weights = np.load(npy_weights).item()
|
||||
def set_npy_weights(self, weights_path):
|
||||
npy_weights_path = weights_path + ".npy"
|
||||
json_path = weights_path + ".json"
|
||||
h5_path = weights_path + ".h5"
|
||||
|
||||
for layer in model.layers:
|
||||
print 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)
|
||||
scale = weights[layer.name]['scale'].reshape(-1)
|
||||
offset = weights[layer.name]['offset'].reshape(-1)
|
||||
|
||||
model.get_layer(layer.name).set_weights([mean, variance, scale, offset])
|
||||
print("Importing weights from %s" % npy_weights_path)
|
||||
weights = np.load(npy_weights_path).item()
|
||||
|
||||
for layer in self.model.layers:
|
||||
print(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)
|
||||
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])
|
||||
|
||||
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:
|
||||
biases = weights[layer.name]['biases']
|
||||
self.model.get_layer(layer.name).set_weights([weight, biases])
|
||||
print('Finished importing weights.')
|
||||
|
||||
print("Writing keras model & weights")
|
||||
json_string = self.model.to_json()
|
||||
with open(json_path, 'w') as file_handle:
|
||||
file_handle.write(json_string)
|
||||
self.model.save_weights(h5_path)
|
||||
print("Finished writing Keras model & weights")
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu':
|
||||
try:
|
||||
weight = weights[layer.name]['weights']
|
||||
model.get_layer(layer.name).set_weights([weight])
|
||||
except Exception as err:
|
||||
biases = weights[layer.name]['biases']
|
||||
model.get_layer(layer.name).set_weights([weight, biases])
|
||||
print 'Finished.'
|
||||
return model
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--input_path', type=str, default='', required=True, help='Path the input image')
|
||||
parser.add_argument('--output_path', type=str, default='', required=True, 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()
|
||||
|
||||
@@ -84,13 +127,27 @@ if __name__ == "__main__":
|
||||
|
||||
with sess.as_default():
|
||||
img = misc.imread(args.input_path)
|
||||
|
||||
pspnet = PSPNet()
|
||||
print(args)
|
||||
|
||||
if "pspnet50" in 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)
|
||||
if "voc2012" in args.model:
|
||||
pspnet = PSPNet101(nb_classes=21, input_shape=(473, 473), weights=args.model)
|
||||
|
||||
else:
|
||||
print("Network architecture not implemented.")
|
||||
|
||||
probs = pspnet.predict(img)
|
||||
print("Writing results...")
|
||||
|
||||
cm = np.argmax(probs, axis=2) + 1
|
||||
pm = np.max(probs, axis=2)
|
||||
color_cm = utils.add_color(cm)
|
||||
misc.imsave(args.output_path, color_cm)
|
||||
misc.imsave("probs.jpg", pm)
|
||||
|
||||
alpha_blended = 0.5 * color_cm * 255 + 0.5 * img # color cm is [0.0-1.0] img [0-255]
|
||||
filename, ext = splitext(args.output_path)
|
||||
misc.imsave(filename + "_seg" + ext, color_cm)
|
||||
misc.imsave(filename + "_probs" + ext, pm)
|
||||
misc.imsave(filename + "_seg_blended" + ext, alpha_blended)
|
||||
|
||||
Reference in New Issue
Block a user