Added script files

README.md

@@ -1 +1,25 @@
-keras
+# Keras implementation of [PSPNet(caffe)](https://github.com/hszhao/PSPNet)
+
+Implemented Architecture of pyramid scene parsing network in Keras
+
+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)
+
+And place in directory with pspnet.py
+
+Memory usage:3500Mb
+Calculation speed: 1.2 sec
+
+## Dependencies:
+1. Tensorflow
+2. Keras
+3. numpy
+
+
+## Usage: 
+
+```bash
+python pspnet.py
+```

drawImage/DrawBbox/DrawBbox.py

@@ -0,0 +1,99 @@
+#OpenCV module for bbox search
+import numpy as np
+import cv2
+from PIL import Image, ImageDraw
+import math
+import random
+import copy
+#JUST WINDOWS  aka binary
+
+#_-------Changed box drawing mode from vertical to with angle ------_#
+class Bbox:
+	def __init__(self):
+		self.bboxInfo = {}
+	def bubble_sort(self, items, numToReturn):
+		""" Implementation of bubble sort """
+		for i in range(len(items)):
+			for j in range(len(items)-1-i):
+					if items[j][1] < items[j+1][1]:
+						items[j], items[j+1] = items[j+1], items[j] 
+		return items[:numToReturn]
+	def filterBboxes(self):
+		for bboxObject in self.objects_to_bbox:
+			if self.class_ratio[bboxObject]<self.bbox_filter:
+				self.bboxInfo.pop(bboxObject)
+
+		
+	def drawLimitedObjects(self, segmented, raw, **kwargs):
+		num_to_draw = kwargs['numToDraw']
+		self.class_ratio = kwargs['classRatio']
+		self.bbox_filter = kwargs['bboxFilter']
+		self.objects_to_bbox = kwargs['classesToBbox']
+		listToDraw = []
+		self.filterBboxes()
+		keys = self.bboxInfo.keys()
+
+		for key in keys:
+			for x in range(self.bboxInfo[key].__len__()):
+				area = self.bboxInfo[key][x][4]
+				listToDraw.append([key, area, self.bboxInfo[key][x]])
+		listToDraw = self.bubble_sort(listToDraw, num_to_draw)
+		print listToDraw
+		segmentedImageRGB = np.array(segmented)
+		output_im = np.array(raw)
+
+		JSONCoords = {}
+		if num_to_draw>listToDraw.__len__():
+			num_to_draw = listToDraw.__len__()
+
+		for i in range(num_to_draw):
+			x1 = listToDraw[i][2][0]
+			y1 = listToDraw[i][2][1]
+			x2 = listToDraw[i][2][0] + listToDraw[i][2][2]
+			y2 = listToDraw[i][2][1] + listToDraw[i][2][3]
+			clr = listToDraw[i][2][5]
+			box = cv2.cv.BoxPoints(((listToDraw[i][2][0],listToDraw[i][2][1]),(listToDraw[i][2][2],listToDraw[i][2][3]),listToDraw[i][2][6])) # cv2.boxPoints(rect) for OpenCV 3.x
+			box = np.int0(box)
+			cv2.drawContours(segmentedImageRGB,[box],0,clr,2)
+			cv2.drawContours(output_im,[box],0,clr,2)
+
+		BboxedRawImage = Image.fromarray(output_im)
+		for x in range(num_to_draw):
+			listToDraw[x][2].pop()
+			listToDraw[x][2].pop()
+			if JSONCoords.has_key(listToDraw[x][0]):
+				JSONCoords[listToDraw[x][0]].append(listToDraw[x][2])
+			else:
+				JSONCoords[listToDraw[x][0]]=[listToDraw[x][2]]
+		BboxedImage = Image.fromarray(segmentedImageRGB)
+		return BboxedImage, BboxedRawImage, JSONCoords #image
+		
+
+
+
+
+	def findBbox(self, segmentedImageRGB, maskImage,  **kwargs):
+		objectName = kwargs['objectName']
+
+		maskImage = maskImage.convert("RGB")
+		imW,imH = segmentedImageRGB.size
+		open_cv_image = np.array(maskImage) 
+		open_cv_image = cv2.cvtColor(open_cv_image, cv2.COLOR_RGB2BGR)
+
+		imgray = cv2.cvtColor(open_cv_image,cv2.COLOR_BGR2GRAY)
+		ret,thresh = cv2.threshold(imgray,127,255,0)
+		contours, _ = cv2.findContours(thresh, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
+		self.bboxInfo[objectName] = [] #[[coords,clr],[coords,clr]]
+		bboxClr = (int(math.floor(random.random()*255)), int(math.floor(random.random()*255)), int(math.floor(random.random()*255)))
+		for c in contours:
+			rect = cv2.minAreaRect(c)
+			if (rect[1][0]*rect[1][1])<500: continue  
+			x = int(rect[0][0])
+			y = int(rect[0][1])
+			w = int(rect[1][0])
+			h = int(rect[1][1])
+			rotation = rect[2]
+			single_bbox_info = [x,y,w,h,w*h,bboxClr, rotation]
+			#Instrument to filter inner bboxes
+			if thresh[y][x]==0: continue
+			self.bboxInfo[objectName].append(single_bbox_info)

drawImage/DrawBbox/__init__.py

@@ -0,0 +1 @@
+from DrawBbox import *

drawImage/__init__.py

@@ -0,0 +1 @@
+from drawModule import *

drawImage/drawModule.py

@@ -0,0 +1,74 @@
+from PIL import Image, ImageDraw
+import scipy.ndimage
+import scipy.io
+import numpy as np
+import time
+import copy
+
+
+class BaseDraw:
+	def __init__(self, color150, objectNames, img, pred_size, predicted_classes):
+		self.class_colors = scipy.io.loadmat(color150)
+		self.class_names = scipy.io.loadmat(objectNames, struct_as_record=False)
+		self.im = img
+		self.pred_size = pred_size
+		self.predicted_classes = copy.deepcopy(predicted_classes)
+		self.original_W = self.im.size[0]
+		self.original_H = self.im.size[1]
+
+		self.output_W = 1920
+		self.output_H = 1080	
+
+
+	def dumpArray(self, array, i):
+		test = array*100
+		test = Image.fromarray(test.astype('uint8'))
+		test = test.convert("RGB")
+		test.save('/home/vlad/oS_AI/'+str(i)+'t.jpg', "JPEG")
+
+	def calculateResize(self):
+		W_coef = float(self.original_W)/float(self.output_W)
+		H_coef = float(self.original_H)/float(self.output_H)
+		horiz_pad = 0
+		vert_pad = 0
+		if W_coef > H_coef:
+			coef = W_coef
+			horiz_pad = int((self.output_H - self.original_H/coef)/2)
+			return [coef, horiz_pad, vert_pad]
+		else:
+			coef = H_coef
+			vert_pad = int((self.output_W - self.original_W/coef)/2)
+			return [coef, horiz_pad, vert_pad]
+
+
+	def resizeToOutput(self, image, coef, h_pad, w_pad):
+		image = image.resize((int(self.original_W/coef), int(self.original_H/coef)), resample=Image.BILINEAR)
+		outputImage = Image.new("RGB",(self.output_W,self.output_H),(0,0,0))
+		outputImage.paste(image,(w_pad,h_pad))
+		return outputImage
+
+
+
+	def drawSimpleSegment(self):
+
+		#Drawing module
+		im_Width, im_Height = self.pred_size
+		prediction_image = Image.new("RGB", (im_Width, im_Height) ,(0,0,0))
+		prediction_imageDraw = ImageDraw.Draw(prediction_image)
+
+		#BASE all image segmentation
+		for i in range(im_Width):
+			for j in range(im_Height):
+				#get matrix element class(0-149)
+				px_Class = self.predicted_classes[j][i]
+				#assign color from .mat list
+				put_Px_Color = tuple(self.class_colors['colors'][px_Class])
+
+				#drawing
+				prediction_imageDraw.point((i,j), fill=put_Px_Color)
+
+		#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)
+
+		return FullHdOutImage

pspnet.py

@@ -0,0 +1,334 @@
+from keras.models import Sequential
+from keras.layers import Conv2D, MaxPooling2D, AveragePooling2D, UpSampling2D
+from keras.layers import BatchNormalization, Activation, Input, Dropout, ZeroPadding2D
+from keras.layers import Add, merge, concatenate, Lambda, Reshape
+from keras import backend as K
+import tensorflow as tf
+from keras.models import Model
+import numpy as np
+from PIL import Image
+import drawImage
+import time
+
+def load_weights():
+	w = np.load('pspnet.npy').item()
+	return w
+
+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)
+
+			offset = weights[layer.name]['offset'].reshape(-1)
+			mean = weights[layer.name]['mean'].reshape(-1)
+			variance = weights[layer.name]['variance'].reshape(-1)
+			
+			model.get_layer(layer.name).set_weights([mean, variance,
+													scale, offset])
+
+		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])
+
+	print 'weights set finish'
+	return model
+
+def Interp_(x, size=None, zoom=None):
+	print(x.shape)
+
+	old_height = int(x.shape[2])
+	old_width = int(x.shape[3])
+	if zoom is not None:
+		zoom = int(zoom)
+		new_height = old_height + (old_height-1) * (zoom - 1)
+		new_width = old_width + (old_width-1) * (old_width - 1)
+	elif size is not None:
+		new_height = size[0]
+		new_width = size[1]
+	resized = tf.image.resize_images(x, [new_height, new_width])
+	return resized
+
+
+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_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
+
+#NOT USED---
+def add_common_layers(prev):
+	prev = BatchNormalization(momentum=0.95)(prev)
+	prev = Activation('relu')(prev)
+	return prev 
+
+def Conv(prev_layer, level, kernel=(1,1), strides=(1,1)):
+	layer = Conv2D(64 * level, (1,1), strides=(1,1))(prev_layer)
+	return layer
+#-----------
+
+
+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 = BatchNormalization(momentum=0.95, 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, use_bias=False,
+					name=names[2])(prev)
+	
+
+	prev = BatchNormalization(momentum=0.95, name=names[3])(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)
+	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])(prev)
+	return prev
+
+
+def empty_branch(prev):
+	return prev
+
+
+def residual_short(prev_layer, level, pad=1, lvl=1, sub_lvl=1, modify_stride=False):
+
+	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') 
+
+
+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_2 = empty_branch(prev_layer)
+	return merge([block_1, block_2], mode='sum') 
+
+def interp_block(prev_layer, level, str_lvl=1):
+
+	str_lvl = str(str_lvl)
+
+	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])(prev_layer)
+	prev_layer = Activation('relu')(prev_layer)
+	prev_layer = Lambda(Interp)(prev_layer)
+	return prev_layer
+
+
+if __name__ == "__main__":
+
+	#Names for the first 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)
+
+	inp = Input((473,473, 3))
+
+	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"
+	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"
+	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"
+	relu1 = Activation('relu')(bn1)				#"conv1_3_3x3/relu"
+
+	res = ZeroPadding2D(padding=(1,1))(relu1)
+	res = MaxPooling2D(pool_size=(3,3), strides=(2,2))(res)  #"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(2): 
+		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")(res)
+	res = Activation('relu')(res)
+	#res = Dropout(0.1)(res)
+	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)
+
+	
+
+
+
+	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('test.jpg')
+		data_im = np.asarray(image)
+		data = np.zeros([1,473,473,3])
+		data_im = np.resize(data_im, [473, 473, 3])
+
+		data[0] = data_im
+
+		#predict
+		
+		startForward = time.time()
+		pred = model.predict(data, batch_size=1, verbose=0)
+		finishForward = (time.time() - startForward)
+
+		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'
+
+
+	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('out.jpg',"JPEG")
+
+