diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..60bf75b --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +*.npy +*.pyc +*.png + diff --git a/README.md b/README.md index 57760da..7b44f28 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,7 @@ Converted trained weights needed to run the network. Download converted weights here: -[link:pspnet50_ade20k.npy](https://www.dropbox.com/s/2ksp9hvokzk6qc8/pspnet50_ade20k.npy?dl=0) +[link:pspnet50_ade20k.npy](https://www.dropbox.com/s/ms8afun494dlh1t/pspnet50_ade20k.npy?dl=0) And place in directory with pspnet50_ade20k.npy @@ -24,6 +24,8 @@ Was repaired some issues. But the result is not as well as expected compared to ![Original](test.jpg) ![Processed](test_seg.jpg) ![Alpha mixed](test_seg_blended.jpg) +![New](out.jpg) +![New](probs.jpg) ## Pycaffe result ![Pycaffe results](test_pycaffe.jpg) diff --git a/layers_builder.py b/layers_builder.py index 5c04427..724f79c 100644 --- a/layers_builder.py +++ b/layers_builder.py @@ -1,225 +1,198 @@ from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D -from keras.layers import BatchNormalization, Activation, Input, Dropout, ZeroPadding2D -from keras.layers import merge, concatenate, Lambda, Reshape +from keras.layers import BatchNormalization, Activation, Input, Dropout, ZeroPadding2D, Lambda +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 +def BN(name=""): + return BatchNormalization(momentum=0.95, name=name, epsilon=1e-5) -def Interp(x, size=(60,60)): - print(x.shape) - new_height = size[0] - new_width = size[1] - resized = tf.image.resize_images(x, [new_height, new_width]) - print(resized.shape) - return resized +def Interp(x, shape=(60,60)): + new_height,new_width = shape + resized = tf.image.resize_images(x, [new_height, new_width], align_corners=True) + return resized +def residual_conv(prev, level, pad=1, lvl=1, sub_lvl=1, modify_stride=False): + lvl = str(lvl) + sub_lvl = str(sub_lvl) + names = ["conv"+lvl+"_"+ sub_lvl +"_1x1_reduce" , + "conv"+lvl+"_"+ sub_lvl +"_1x1_reduce_bn", + "conv"+lvl+"_"+ sub_lvl +"_3x3", + "conv"+lvl+"_"+ sub_lvl +"_3x3_bn", + "conv"+lvl+"_"+ sub_lvl +"_1x1_increase", + "conv"+lvl+"_"+ sub_lvl +"_1x1_increase_bn"] + if modify_stride == False: + prev = Conv2D(64 * level, (1,1), strides=(1,1), name=names[0], use_bias=False)(prev) + elif modify_stride == True: + prev = Conv2D(64 * level, (1,1), strides=(2,2), name=names[0], use_bias=False)(prev) -def Interp_zoom(x, zoom=8): - print(x.shape) - old_height = int(x.shape[1]) - old_width = int(x.shape[2]) - new_height = old_height + (old_height-1) * (zoom - 1) - new_width = old_width + (old_width-1) * (zoom - 1) - resized = tf.image.resize_images(x, [new_height, new_width]) - return resized + 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) -def residual_conv(prev, level, - pad=1, lvl=1, sub_lvl=1, modify_stride=False): - - lvl = str(lvl) - sub_lvl = str(sub_lvl) - names = ["conv"+lvl+"_"+ sub_lvl +"_1x1_reduce" , - "conv"+lvl+"_"+ sub_lvl +"_1x1_reduce_bn", - "conv"+lvl+"_"+ sub_lvl +"_3x3", - "conv"+lvl+"_"+ sub_lvl +"_3x3_bn", - "conv"+lvl+"_"+ sub_lvl +"_1x1_increase", - "conv"+lvl+"_"+ sub_lvl +"_1x1_increase_bn"] - if modify_stride == False: - prev = Conv2D(64 * level, (1,1), strides=(1,1), use_bias=False, - name=names[0])(prev) - elif modify_stride == True: - prev = Conv2D(64 * level, (1,1), strides=(2,2), use_bias=False, - name=names[0])(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 = BN(name=names[5])(prev) + return prev - prev = BatchNormalization(momentum=0.95, name=names[1], epsilon=1e-5)(prev) - prev = Activation('relu')(prev) +def short_convolution_branch(prev, level, lvl=1, sub_lvl=1, modify_stride=False): + lvl = str(lvl) + sub_lvl = str(sub_lvl) + names = ["conv"+lvl+"_"+ sub_lvl +"_1x1_proj", + "conv"+lvl+"_"+ sub_lvl +"_1x1_proj_bn"] - prev = ZeroPadding2D(padding=(pad,pad))(prev) - prev = Conv2D(64 * level, (3,3), - strides=(1,1), dilation_rate=pad, use_bias=False, - name=names[2])(prev) - - - prev = BatchNormalization(momentum=0.95, name=names[3], epsilon=1e-5)(prev) - prev = Activation('relu')(prev) - prev = Conv2D(256 * level, (1,1), strides=(1,1), use_bias=False, - name=names[4])(prev) - prev = BatchNormalization(momentum=0.95, name=names[5], epsilon=1e-5)(prev) - return prev - - -def short_convolution_branch(prev, level, - lvl=1, sub_lvl=1, modify_stride=False): - lvl = str(lvl) - sub_lvl = str(sub_lvl) - names = ["conv"+lvl+"_"+ sub_lvl +"_1x1_proj", - "conv"+lvl+"_"+ sub_lvl +"_1x1_proj_bn" - ] - - if modify_stride == False: - prev = Conv2D(256 * level ,(1,1), strides=(1,1), use_bias=False, - name=names[0])(prev) - elif modify_stride == True: - prev = Conv2D(256 * level, (1,1), strides=(2,2), use_bias=False, - name=names[0])(prev) - - prev = BatchNormalization(momentum=0.95, name=names[1], epsilon=1e-5)(prev) - return prev + if modify_stride == False: + prev = Conv2D(256 * level ,(1,1), strides=(1,1), name=names[0], use_bias=False)(prev) + elif modify_stride == True: + prev = Conv2D(256 * level, (1,1), strides=(2,2), name=names[0], use_bias=False)(prev) + prev = BN(name=names[1])(prev) + return prev def empty_branch(prev): - return prev - + return prev def residual_short(prev_layer, level, pad=1, lvl=1, sub_lvl=1, modify_stride=False): - prev_layer = Activation('relu')(prev_layer) - block_1 = residual_conv(prev_layer, level, - pad=pad, lvl=lvl, sub_lvl=sub_lvl, - modify_stride=modify_stride) - - block_2 = short_convolution_branch(prev_layer, level, - lvl=lvl, sub_lvl=sub_lvl, - modify_stride=modify_stride) - - return merge([block_1, block_2], mode='sum') + prev_layer = Activation('relu')(prev_layer) + block_1 = residual_conv(prev_layer, level, + pad=pad, lvl=lvl, sub_lvl=sub_lvl, + modify_stride=modify_stride) + block_2 = short_convolution_branch(prev_layer, level, + lvl=lvl, sub_lvl=sub_lvl, + modify_stride=modify_stride) + added = Add()([block_1, block_2]) + return added def residual_empty(prev_layer, level, pad=1, lvl=1, sub_lvl=1): - prev_layer = Activation('relu')(prev_layer) + prev_layer = Activation('relu')(prev_layer) - block_1 = residual_conv(prev_layer, level, - pad=pad, lvl=lvl, sub_lvl=sub_lvl) - block_2 = empty_branch(prev_layer) - return merge([block_1, block_2], mode='sum') + 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 +def ResNet(inp): + #Names for the first couple layers of model + names = ["conv1_1_3x3_s2", + "conv1_1_3x3_s2_bn", + "conv1_2_3x3", + "conv1_2_3x3_bn", + "conv1_3_3x3", + "conv1_3_3x3_bn"] + + #---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" + 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" + 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" + 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" + + #---Residual layers(body of network) + + """ + Modify_stride --Used only once in first 3_1 convolutions block. + changes stride of first convolution from 1 -> 2 + """ + + #2_1- 2_3 + res = residual_short(res, 1, pad=1, lvl=2, sub_lvl=1) + for i in range(2): + res = residual_empty(res, 1, pad=1, lvl=2, sub_lvl=i+2) + + #3_1 - 3_3 + res = residual_short(res, 2, pad=1, lvl=3, sub_lvl=1, modify_stride=True) + for i in range(3): + res = residual_empty(res, 2, pad=1, lvl=3, sub_lvl=i+2) + + #4_1 - 4_6 + res = residual_short(res, 4, pad=2, lvl=4, sub_lvl=1) + for i in range(5): + res = residual_empty(res, 4, pad=2, lvl=4, sub_lvl=i+2) + + #5_1 - 5_3 + res = residual_short(res, 8, pad=4, lvl=5, sub_lvl=1) + for i in range(2): + res = residual_empty(res, 8, pad=4, lvl=5, sub_lvl=i+2) + + res = Activation('relu')(res) + return res def interp_block(prev_layer, level, str_lvl=1): - str_lvl = str(str_lvl) + str_lvl = str(str_lvl) - names = [ - "conv5_3_pool"+str_lvl+"_conv", - "conv5_3_pool"+str_lvl+"_conv_bn" - ] + names = [ + "conv5_3_pool"+str_lvl+"_conv", + "conv5_3_pool"+str_lvl+"_conv_bn" + ] - 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), use_bias=False, name=names[0])(prev_layer) - prev_layer = BatchNormalization(momentum=0.95, name=names[1], epsilon=1e-5)(prev_layer) - prev_layer = Activation('relu')(prev_layer) - prev_layer = Lambda(Interp)(prev_layer) - return prev_layer + 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 = BN(name=names[1])(prev_layer) + prev_layer = Activation('relu')(prev_layer) + prev_layer = Lambda(Interp)(prev_layer) + return prev_layer +def PSPNet(res): -def build_pspnet(): - #Names for the first couple layers of model - names = ["conv1_1_3x3_s2", - "conv1_1_3x3_s2_bn", - "conv1_2_3x3", - "conv1_2_3x3_bn", - "conv1_3_3x3", - "conv1_3_3x3_bn"] + #---PSPNet concat layers with Interpolation - #---Short branch(only start of network) + interp_block1 = interp_block(res, 6, str_lvl=1) + interp_block2 = interp_block(res, 3, str_lvl=2) + interp_block3 = interp_block(res, 2, str_lvl=3) + interp_block6 = interp_block(res, 1, str_lvl=6) - inp = Input((473,473, 3)) + #concat all these layers. resulted shape=(1,60,60,4096) + res = Concatenate()([res, + interp_block6, + interp_block3, + interp_block2, + interp_block1]) + return res - cnv1 = ZeroPadding2D(padding=(1,1))(inp) - cnv1 = Conv2D(64, (3, 3), strides=(2, 2), use_bias=False, name=names[0])(cnv1) # "conv1_1_3x3_s2" - - bn1 = BatchNormalization(momentum=0.95, name=names[1], epsilon=1e-5)(cnv1) # "conv1_1_3x3_s2/bn" - relu1 = Activation('relu')(bn1) #"conv1_1_3x3_s2/relu" +def build_pspnet(activation='softmax'): + ''' + Normal PSPNet. + ''' + inp = Input((473,473,3)) + res = ResNet(inp) + psp = PSPNet(res) - cnv1 = ZeroPadding2D(padding=(1,1))(relu1) - cnv1 = Conv2D(64, (3, 3), strides=(1, 1), use_bias=False, name=names[2])(cnv1) #"conv1_2_3x3" - - bn1 = BatchNormalization(momentum=0.95, name=names[3], epsilon=1e-5)(cnv1) #"conv1_2_3x3/bn" - relu1 = Activation('relu')(bn1) #"conv1_2_3x3/relu" + 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) - cnv1 = ZeroPadding2D(padding=(1,1))(relu1) - cnv1 = Conv2D(128, (3, 3), strides=(1, 1), use_bias=False, name=names[4])(cnv1) #"conv1_3_3x3" - - bn1 = BatchNormalization(momentum=0.95, name=names[5], epsilon=1e-5)(cnv1) #"conv1_3_3x3/bn" - relu1 = Activation('relu')(bn1) #"conv1_3_3x3/relu" + x = Conv2D(150, (1, 1), strides=(1, 1), name="conv6")(x) + x = Lambda(Interp, arguments={'shape': (473,473)})(x) + x = Activation('softmax')(x) - res = ZeroPadding2D(padding=(1,1))(relu1) - res = MaxPooling2D(pool_size=(3,3), strides=(2,2))(res) #"pool1_3x3_s2" - + model = Model(inputs=inp, outputs=x) - #---Residual layers(body of network) - - """ - Modify_stride --Used only once in first 3_1 convolutions block. - changes stride of first convolution from 1 -> 2 - """ - - #2_1- 2_3 - res = residual_short(res, 1, pad=1, lvl=2, sub_lvl=1) - for i in range(2): - res = residual_empty(res, 1, pad=1, lvl=2, sub_lvl=i+2) - - #3_1 - 3_3 - res = residual_short(res, 2, pad=1, lvl=3, sub_lvl=1, modify_stride=True) - for i in range(3): #for i in range(2): old wrong code - res = residual_empty(res, 2, pad=1, lvl=3, sub_lvl=i+2) - - #4_1 - 4_6 - res = residual_short(res, 4, pad=2, lvl=4, sub_lvl=1) - for i in range(5): - res = residual_empty(res, 4, pad=2, lvl=4, sub_lvl=i+2) - - #5_1 - 5_3 - res = residual_short(res, 8, pad=4, lvl=5, sub_lvl=1) - for i in range(2): - res = residual_empty(res, 8, pad=4, lvl=5, sub_lvl=i+2) - - #---Head of network - #---PSPNet concat layers with Interpolation - - res = Activation('relu')(res) - interp_block1 = interp_block(res, 6, str_lvl=1) - interp_block2 = interp_block(res, 3, str_lvl=2) - interp_block3 = interp_block(res, 2, str_lvl=3) - interp_block4 = interp_block(res, 1, str_lvl=6) - - #concat all these layers by 4th axis(3+1). resulted shape=(1,60,60,4096) - res = concatenate([res, - interp_block1, - interp_block2, - interp_block3, - interp_block4], axis=3) - - res = ZeroPadding2D(padding=(1,1))(res) - res = Conv2D(512, (3, 3), strides=(1, 1), use_bias=False, name="conv5_4")(res) - - res = BatchNormalization(momentum=0.95, name="conv5_4_bn", epsilon=1e-5)(res) - res = Activation('relu')(res) - #res = Dropout(0.1)(res) #used only in training - res = Conv2D(150, (1, 1), strides=(1, 1), name="conv6")(res) - res = Lambda(Interp_zoom)(res) - - - #Use softmax layer for pixelwise prediction - curr_width, curr_height, curr_channels = res._shape_as_list()[1:] - - reshape = Reshape((curr_width*curr_height, curr_channels))(res) - activation = Activation('softmax')(reshape) - reshape = Reshape((curr_width, curr_height, curr_channels))(activation) - - #End of model - model = Model(inputs=inp, outputs=reshape) - return model + # Solver + sgd = SGD(lr=learning_rate, momentum=0.9, nesterov=True) + model.compile(optimizer=sgd, + loss='categorical_crossentropy', + metrics=['accuracy']) + return model diff --git a/out.jpg b/out.jpg new file mode 100644 index 0000000..14c2689 Binary files /dev/null and b/out.jpg differ diff --git a/probs.jpg b/probs.jpg new file mode 100644 index 0000000..9bc9563 Binary files /dev/null and b/probs.jpg differ diff --git a/pspnet.py b/pspnet.py index 2374e9a..fa42f64 100644 --- a/pspnet.py +++ b/pspnet.py @@ -1,112 +1,96 @@ -from keras import backend as K -from PIL import Image - -import layers_builder as pspnet -import tensorflow as tf -import numpy as np -import drawImage +import os import argparse -import time +import numpy as np +from scipy import misc, ndimage +from keras import backend as K +import tensorflow as tf +import layers_builder as layers +import utils -def load_weights(): - w = np.load('pspnet50_ade20k.npy').item() - return w +WEIGHTS = 'pspnet50_ade20k.npy' +DATA_MEAN = np.array([[[123.68, 116.779, 103.939]]]) # RGB +class PSPNet: -def set_weights(model, weights): - print 'weights set start' - for layer in model.layers: - if layer.name[:4] == 'conv' and layer.name[-2:] == 'bn': - print layer.name - scale = weights[layer.name]['scale'].reshape(-1) + def __init__(self): + self.model = layers.build_pspnet() + set_npy_weights(self.model, WEIGHTS) - offset = weights[layer.name]['offset'].reshape(-1) - mean = weights[layer.name]['mean'].reshape(-1) - variance = weights[layer.name]['variance'].reshape(-1) + def predict(self, img): + ''' + Arguments: + img: must be 473x473x3 + ''' + h_ori,w_ori = img.shape[:2] - # mean *= scale - # variance *= scale - - # model.get_layer(layer.name).set_weights([mean, variance, - # scale, offset]) - model.get_layer(layer.name).set_weights([scale, offset, - mean, variance]) - # model.get_layer(layer.name).set_weights([scale, offset, - # mean, variance]) + # Preprocess + img = misc.imresize(img, (473, 473)) + img = img - DATA_MEAN + img = img[:,:,::-1] # RGB => BGR + img = img.astype('float32') - elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu': - print layer.name - 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]) + 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) + return probs - print 'weights set finish' - return model + def predict_sliding_window(self, img): + pass + def feed_forward(self, data): + assert data.shape == (473,473,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() + + 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]) + + 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('--id', default="0") + args = parser.parse_args() - settings = None - 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') + os.environ["CUDA_VISIBLE_DEVICES"] = args.id - settings, unparsed = parser.parse_known_args() - mean_r = 123.68 - mean_g = 116.779 - mean_b = 103.939 + sess = tf.Session() + K.set_session(sess) - model = pspnet.build_pspnet() - - sess = tf.Session() - K.set_session(sess) - - with sess.as_default(): - #Load weights into variable - npy_weights = load_weights() - #Set weights to each laye by name - model = set_weights(model, npy_weights) - - #Load image, resize and paste into 4D tensor - image = Image.open(settings.input_path) - im = image.resize((473, 473)) - input_ = np.array(im, dtype=np.float32) - input_ = input_[:,:,::-1] - input_ -= np.array((mean_b, mean_g, mean_r)) - data = np.zeros([1,473,473,3]) - - data[0] = input_ - - #predict - - startForward = time.time() - pred = model.predict(data, batch_size=1, verbose=0) - finishForward = (time.time() - startForward) - print "Time used: %f" % finishForward - # pred = np.transpose(pred[0], (2, 1, 0)) - print np.shape(pred) - pred = pred[0] - predicted_classes = np.argmax(pred, axis=2) - - proto = 'utils/model/pspnet.prototxt' - weights = 'utils/model/pspnet.caffemodel' - colors = 'utils/colorization/color150.mat' - objects = 'utils/colorization/objectName150.mat' - - - im_Width = predicted_classes.shape[0] - im_Height = predicted_classes.shape[1] - draw = drawImage.BaseDraw(colors, objects, - image, (im_Width, im_Height), - predicted_classes) - simpleSegmentImage = draw.drawSimpleSegment(); - simpleSegmentImage.save(settings.output_path,"JPEG") + with sess.as_default(): + img = misc.imread(args.input_path) + + pspnet = PSPNet() + probs = pspnet.predict(img) + 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) diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..bbf9d82 --- /dev/null +++ b/utils.py @@ -0,0 +1,32 @@ +import colorsys +import numpy as np + +from keras.models import Model + +def add_color(img): + h,w = img.shape + img_color = np.zeros((h,w,3)) + for i in xrange(1,151): + img_color[img == i] = to_color(i) + return img_color + +def to_color(category): + # Maps each category a good distance away + # from each other on the HSV color space + v = (category-1)*(137.5/360) + return colorsys.hsv_to_rgb(v,1,1) + + +# For printing the activations in each layer +# Useful for debugging +def debug(model, data): + names = [layer.name for layer in model.layers] + for name in names[:]: + print_activation(model, name, data) +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) +def array_to_str(a): + return "{} {} {} {} {}".format(a.dtype, a.shape, np.min(a), np.max(a), np.mean(a)) \ No newline at end of file diff --git a/weight_converter.py b/weight_converter.py index c97241f..033beb1 100644 --- a/weight_converter.py +++ b/weight_converter.py @@ -1,20 +1,42 @@ -import caffe +import os +import sys import numpy as np -import os, sys + +import caffe + +# Not needed because Tensorflow and Caffe do convolution the same way +# Needed for conversion to Theano +def rot90(W): + for i in range(W.shape[0]): + for j in range(W.shape[1]): + W[i, j] = np.rot90(W[i, j], 2) + return W weights = {} net = caffe.Net(sys.argv[1], sys.argv[2], caffe.TEST) for k,v in net.params.items(): print "Layer %s, has %d params." % (k, len(v)) if len(v) == 1: - weights[k] = {"weights": np.transpose(v[0].data[...], (2,3,1,0))} + W = v[0].data[...] + W = np.transpose(W, (2,3,1,0)) + weights[k] = {"weights": W} elif len(v) == 2: - weights[k] = {"weights": np.transpose(v[0].data[...], (2,3,1,0)), "biases": v[1].data[...]} + W = v[0].data[...] + W = np.transpose(W, (2,3,1,0)) + b = v[1].data[...] + weights[k] = {"weights": W, "biases": b} elif len(v) == 4: - weights[k.replace('/', '_')] = {"scale": v[0].data[...], "offset": v[1].data[...], "mean": v[2].data[...], "variance": v[3].data[...]} + k = k.replace('/', '_') + mean = v[0].data[...] + variance = v[1].data[...] + scale = v[2].data[...] + offset = v[3].data[...] + weights[k] = {"mean": mean, "variance": variance, "scale": scale, "offset": offset} else: print "Undefined layer" exit() arr = np.asarray(weights) -np.save("pspnet50_ade20k.npy", arr) \ No newline at end of file +np.save("pspnet50_ade20k.npy", arr) + +