diff --git a/README.md b/README.md
index 96756b9..468ff93 100644
--- a/README.md
+++ b/README.md
@@ -6,26 +6,30 @@ Converted trained weights needed to run the network.
 
 
 Download converted weights here:
-[link:pspnet.npy](https://www.dropbox.com/s/9xebhix7dbk372d/pspnet.npy?dl=0)
+[link:pspnet50_ade20k.npy](https://www.dropbox.com/s/2ksp9hvokzk6qc8/pspnet50_ade20k.npy?dl=0)
 
-And place in directory with pspnet.py
+And place in directory with pspnet50_ade20k.npy
 
-Weights from caffemodel were converted with [caffe-tensorflow](https://github.com/ethereon/caffe-tensorflow), source code of converter was modified to fit batch normalization, which is annotated as 'BN' in original prototxt
+Weights from caffemodel were converted by, weight_converter.py. The usage of this file is
+```bash
+python weight_converter.py <path to .prototxt> <path to .caffemodel>
+```
+Running this need to compile the original PSPNet caffe code and pycaffe. 
 
 Interpolation layer is implemented in code as custom layer "Interp"
 
 ## Important
 
-This implementation is not working properly despite calculations are made without errors(output image is very very bad).
-I can't figure out which causes such behavior, so help and proposals are appreciated.
-
-Memory usage:3500Mb
-Calculation speed: 1.2 sec on gtx 1080
+This repo is forked from [Vladkryvoruchko/PSPNet-Keras-tensorflow](https://github.com/Vladkryvoruchko/PSPNet-Keras-tensorflow) and repaired some issues. But the result is not as well as expected. Here is an example of current code:
+![Original](test.jpg)
+![Processed](test_seg.jpg)
+![Alpha mixed](test_seg_blended.jpg)
 
 ## Dependencies:
 1. Tensorflow
 2. Keras
 3. numpy
+4. pycaffe(PSPNet)(optional)
 
 
 ## Usage: 
diff --git a/drawImage/__init__.pyc b/drawImage/__init__.pyc
new file mode 100644
index 0000000..80696bb
Binary files /dev/null and b/drawImage/__init__.pyc differ
diff --git a/drawImage/drawModule.py b/drawImage/drawModule.py
index e591ca4..42d0c6a 100644
--- a/drawImage/drawModule.py
+++ b/drawImage/drawModule.py
@@ -16,8 +16,8 @@ class BaseDraw:
 		self.original_W = self.im.size[0]
 		self.original_H = self.im.size[1]
 
-		self.output_W = 1920
-		self.output_H = 1080	
+		self.output_W = self.original_W
+		self.output_H = self.original_H	
 
 
 	def dumpArray(self, array, i):
@@ -70,5 +70,6 @@ class BaseDraw:
 		#Resize to original size and save
 		self.coef, self.h_pad, self.w_pad = self.calculateResize()
 		FullHdOutImage = self.resizeToOutput(prediction_image, self.coef, self.h_pad, self.w_pad)
+		FullHdOutImage = Image.blend(FullHdOutImage, self.im, 0.5)
 
 		return FullHdOutImage
diff --git a/drawImage/drawModule.pyc b/drawImage/drawModule.pyc
new file mode 100644
index 0000000..5a8bc97
Binary files /dev/null and b/drawImage/drawModule.pyc differ
diff --git a/layers_builder.py b/layers_builder.py
index 2f8ff16..5c04427 100644
--- a/layers_builder.py
+++ b/layers_builder.py
@@ -44,7 +44,7 @@ def residual_conv(prev, level,
 		prev = Conv2D(64 * level, (1,1), strides=(2,2), use_bias=False,
 					name=names[0])(prev)
 
-	prev = BatchNormalization(momentum=0.95, name=names[1])(prev)
+	prev = BatchNormalization(momentum=0.95, name=names[1], epsilon=1e-5)(prev)
 	prev = Activation('relu')(prev)
 
 	prev = ZeroPadding2D(padding=(pad,pad))(prev)
@@ -53,11 +53,11 @@ def residual_conv(prev, level,
 					name=names[2])(prev)
 	
 
-	prev = BatchNormalization(momentum=0.95, name=names[3])(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])(prev)
+	prev = BatchNormalization(momentum=0.95, name=names[5], epsilon=1e-5)(prev)
 	return prev
 
 
@@ -76,7 +76,7 @@ def short_convolution_branch(prev, level,
 		prev = Conv2D(256 * level, (1,1), strides=(2,2), use_bias=False,
 				name=names[0])(prev)
 
-	prev = BatchNormalization(momentum=0.95, name=names[1])(prev)
+	prev = BatchNormalization(momentum=0.95, name=names[1], epsilon=1e-5)(prev)
 	return prev
 
 
@@ -85,7 +85,7 @@ def empty_branch(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)
@@ -119,7 +119,7 @@ def interp_block(prev_layer, level, str_lvl=1):
 	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])(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
@@ -141,19 +141,19 @@ def build_pspnet():
 	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])(cnv1)  # "conv1_1_3x3_s2/bn"
+	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"
 
 	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])(cnv1)  #"conv1_2_3x3/bn"
+	bn1 = BatchNormalization(momentum=0.95, name=names[3], epsilon=1e-5)(cnv1)  #"conv1_2_3x3/bn"
 	relu1 = Activation('relu')(bn1) 				#"conv1_2_3x3/relu"
 
 	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])(cnv1) 		#"conv1_3_3x3/bn"
+	bn1 = BatchNormalization(momentum=0.95, name=names[5], epsilon=1e-5)(cnv1) 		#"conv1_3_3x3/bn"
 	relu1 = Activation('relu')(bn1)				#"conv1_3_3x3/relu"
 
 	res = ZeroPadding2D(padding=(1,1))(relu1)
@@ -174,7 +174,7 @@ def build_pspnet():
 
 	#3_1 - 3_3
 	res = residual_short(res, 2, pad=1, lvl=3, sub_lvl=1, modify_stride=True) 
-	for i in range(2): 
+	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
@@ -206,7 +206,7 @@ def build_pspnet():
 	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")(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)
diff --git a/pspnet.py b/pspnet.py
index 117bda2..2374e9a 100644
--- a/pspnet.py
+++ b/pspnet.py
@@ -11,7 +11,7 @@ import time
 
 
 def load_weights():
-	w = np.load('pspnet.npy').item()
+	w = np.load('pspnet50_ade20k.npy').item()
 	return w
 
 
@@ -25,9 +25,16 @@ def set_weights(model, weights):
 			offset = weights[layer.name]['offset'].reshape(-1)
 			mean = weights[layer.name]['mean'].reshape(-1)
 			variance = weights[layer.name]['variance'].reshape(-1)
+
+			# mean *= scale
+			# variance *= scale
 			
-			model.get_layer(layer.name).set_weights([mean, variance,
-													scale, offset])
+			# 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])
 
 		elif layer.name[:4] == 'conv' and not layer.name[-4:] == 'relu':
 			print layer.name
@@ -52,7 +59,9 @@ if __name__ == "__main__":
 						required=True, help='Path to output')
 
 	settings, unparsed = parser.parse_known_args()
-
+	mean_r = 123.68
+	mean_g = 116.779
+	mean_b = 103.939
 
 	model = pspnet.build_pspnet()
 
@@ -67,34 +76,37 @@ if __name__ == "__main__":
 
 		#Load image, resize and paste into 4D tensor
 		image = Image.open(settings.input_path)
-		data_im = np.asarray(image)
+		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_im = np.resize(data_im, [473, 473, 3])
 
-		data[0] = data_im
+		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)
 
-		pred = np.transpose(pred[0], (2, 1, 0))
-		predicted_classes = np.argmax(pred, axis=0)
+		proto = 'utils/model/pspnet.prototxt'
+		weights = 'utils/model/pspnet.caffemodel'
+		colors = 'utils/colorization/color150.mat'
+		objects = 'utils/colorization/objectName150.mat'
 
 
-	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[1]
-	im_Height = predicted_classes.shape[0]
-	draw = drawImage.BaseDraw(colors, objects,
-						image, (im_Width, im_Height),
-						predicted_classes)
-	simpleSegmentImage = draw.drawSimpleSegment();
-	simpleSegmentImage.save(settings.output_path,"JPEG")
+		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")
 
 
diff --git a/test_seg.jpg b/test_seg.jpg
new file mode 100644
index 0000000..4ef7c22
Binary files /dev/null and b/test_seg.jpg differ
diff --git a/test_seg_blended.jpg b/test_seg_blended.jpg
new file mode 100644
index 0000000..05d5591
Binary files /dev/null and b/test_seg_blended.jpg differ
diff --git a/weight_converter.py b/weight_converter.py
new file mode 100644
index 0000000..c97241f
--- /dev/null
+++ b/weight_converter.py
@@ -0,0 +1,20 @@
+import caffe
+import numpy as np
+import os, sys
+
+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))}
+	elif len(v) == 2:
+		weights[k] = {"weights": np.transpose(v[0].data[...], (2,3,1,0)), "biases": v[1].data[...]}
+	elif len(v) == 4:
+		weights[k.replace('/', '_')] = {"scale": v[0].data[...], "offset": v[1].data[...], "mean": v[2].data[...], "variance": v[3].data[...]}
+	else:
+		print "Undefined layer"
+		exit()
+
+arr = np.asarray(weights)
+np.save("pspnet50_ade20k.npy", arr)
\ No newline at end of file