initial commit

Data_helper.py

@@ -0,0 +1,151 @@
+from Drawer import *
+import cPickle as pickle
+from copy import deepcopy
+import subprocess
+mult = 1.
+
+def shuffle(train_path, file, expectC, S, batch, epoch):
+	with open(file,'rb') as f:
+		pick, data = pickle.load(f)
+	C = len(pick)
+	if C != expectC:
+		exit("There is a mismatch between the model and the parsed annotations")
+	size = len(data)
+	print 'Dataset of total {}'.format(size)
+	batch_per_epoch = int(size / batch)
+
+	for i in range(epoch):
+		print 'EPOCH {}'.format(i+1)
+		# Shuffle data
+		shuffle_idx = np.random.permutation(np.arange(size))
+		for b in range(batch_per_epoch):
+			for r in range(1):
+				start_idx = b * batch
+				end_idx = (b+1) * batch
+
+				datum = list()
+				x_batch = list()
+				jpgs = list()
+				try:
+				# if True:
+					for j in range(start_idx,end_idx):
+						real_idx = shuffle_idx[j]
+						this = data[real_idx]
+						jpg = this[0]
+						w, h, allobj_ = this[1]
+						allobj = deepcopy(allobj_)
+						flip = (r / 2)  + (r % 2) * (j % 2)
+						flip = flip % 2
+
+						path = '{}{}'.format(train_path, jpg)
+						img, allobj = crop(path, allobj)
+
+						if flip == 1: 
+							img = img[:,:,::-1,:]
+
+						img = [img]
+						jpgs += [path]
+
+						cellx = 1. * w / S
+						celly = 1. * h / S
+						for x in allobj:
+							# cv2.rectangle(img[0], (x[1], x[2]), (x[3], x[4]), (0,0,255), 2)
+							centerx = .5*(x[1]+x[3]) #xmin, xmax
+							centery = .5*(x[2]+x[4]) #ymin, ymax
+							if flip == 1:
+								centerx = w - centerx
+							cx = centerx / cellx
+							cy = centery / celly
+							x[3] = float(x[3]-x[1]) / w
+							x[4] = float(x[4]-x[2]) / h
+							x[3] = np.sqrt(x[3])
+							x[4] = np.sqrt(x[4])
+							x[1] = cx - np.floor(cx)
+							x[2] = cy - np.floor(cy)
+							x += [np.floor(cx)] 
+							x += [np.floor(cy)]
+
+						if False:
+							for x in allobj:
+								cx = x[5] + x[1]
+								cy = x[6] + x[2]
+								centerx = cx * cellx
+								centery = cy * celly
+								ww = x[3] * x[3] * w
+								hh = x[4] * x[4] * h
+								cv2.rectangle(im,
+									(int(centerx - ww/2), int(centery - hh/2)),
+									(int(centerx + ww/2), int(centery + hh/2)),
+									(0,0,255), 2)
+
+							cv2.imshow("result", im)
+							cv2.waitKey()
+							cv2.destroyAllWindows()
+
+
+						probs = np.zeros([S*S,C])
+						confs = np.zeros([S*S,2])
+						coord = np.zeros([S*S,2,4])
+						proid = np.zeros([S*S,C])
+						conid = np.zeros([S*S,2])
+						cooid1 = cooid2 = np.zeros([S*S,1,4])
+						prear = np.zeros([S*S,4])
+						for x in allobj:
+							at = int(x[6] * S + x[5])
+							probs[at, :] = [0.] * C
+							probs[at, pick.index(x[0])] = 1.
+							proid[at, :] = [1] * C
+							coord[at, 0, :] = x[1:5]
+							coord[at, 1, :] = x[1:5]
+							prear[at,0] = x[1] - x[3]**2 * 3.5 # xleft
+							prear[at,1] = x[2] - x[4]**2 * 3.5 # yup
+							prear[at,2] = x[1] + x[3]**2 * 3.5 # xright
+							prear[at,3] = x[2] + x[4]**2 * 3.5 # ybot
+							confs[at, :] = [1.] * 2
+							conid[at, :] = [1.] * 2
+							cooid1[at, 0, :] = [1.] * 4
+							cooid2[at, 0, :] = [1.] * 4
+						upleft   = np.expand_dims(prear[:,0:2], 1) # 49 x 1 
+						botright = np.expand_dims(prear[:,2:4], 1)
+					#==================================================
+						probs = probs.reshape([-1]) # true_class
+						confs1 = confs[:,0]
+						confs2 = confs[:,1]
+						coord = coord.reshape([-1]) # true_coo
+						upleft   = np.concatenate([upleft]*2,1)
+						botright = np.concatenate([botright]*2,1)
+						proid = proid.reshape([-1]) # class_idtf
+						conid1 = conid[:,0]
+						conid2 = conid[:,1]
+						cooid1 = cooid1
+						cooid2 = cooid2
+					#==================================================
+						new = [
+							[probs], [confs1], [confs2], [coord],
+							[upleft], [botright],
+							[proid], [conid1], [conid2], [cooid1], [cooid2]
+						]
+						if datum == list():
+							datum = new
+							x_batch = img
+						else:
+							x_batch += img
+							for i in range(len(datum)):
+								datum[i] = np.concatenate([datum[i], new[i]])
+
+						if False:
+							here = 0
+							names = list()
+							while here + C < S*S*C:
+								consider = probs[here:here+C]
+								if (np.sum(consider) > 0.5):
+									names += [pick[np.argmax(consider)]]
+								here += C
+							print '{} : {}'.format(jpg, names)
+
+
+					x_batch = np.concatenate(x_batch, 0)
+					yield (x_batch, datum)
+				except:
+					print 'Random scale/translate sends object(s) out of bound'
+					continue

Drawer.py

@@ -0,0 +1,130 @@
+from box import *
+from PIL import Image, ImageFile
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+import cv2
+
+def fix(x,c):
+	return max(min(x,c),0)
+
+def crop(imPath, allobj = None):
+
+	im = cv2.imread(imPath)
+	if allobj is not None:
+		h, w, _ = im.shape
+		scale = np.random.uniform()/3. + 1.
+		max_offx = (scale-1.) * w
+		max_offy = (scale-1.) * h
+		offx = int(np.random.uniform() * max_offx)
+		offy = int(np.random.uniform() * max_offy)
+		im = cv2.resize(im, (0,0), fx = scale, fy = scale)
+		im = im[offy : (offy + h), offx : (offx + w)]
+		#---------------
+		# (x,y) --> (scale*x, scale*y)
+		# (scale*x - offx, scale*y - offy)
+		#--------------
+		for obj in allobj:
+			obj[1] = int(obj[1]*scale-offx)
+			obj[3] = int(obj[3]*scale-offx)
+			obj[2] = int(obj[2]*scale-offy)
+			obj[4] = int(obj[4]*scale-offy)
+			obj[1] = fix(obj[1], w)
+			obj[3] = fix(obj[3], w)
+			obj[2] = fix(obj[2], h)
+			obj[4] = fix(obj[4], h)
+			#print obj, w, h
+
+	# return im
+	im_ = cv2.resize(im, (448, 448))
+	image_array = np.array(im_)
+	image_array = image_array / 255.
+	image_array = image_array * 2. - 1.
+	image_array = np.expand_dims(image_array, 0) # 1, height, width, 3
+
+	if allobj is not None:
+		return image_array, allobj
+	else:
+		return image_array
+
+def to_color(indx, base):
+	base2 = base * base
+	b = indx / base2
+	r = (indx % base2) / base
+	g = (indx % base2) % base
+	return (b * 127, r * 127, g * 127)
+
+def draw_predictions(predictions, 
+	img_path, flip, threshold,
+	C, S, labels, colors):
+
+	B = 2
+	boxes = []
+	SS        =  S * S # number of grid cells
+	prob_size = SS * C # class probabilities
+	conf_size = SS * B # confidences for each grid cell
+	probs = predictions[0 : prob_size]
+	confs = predictions[prob_size : (prob_size + conf_size)]
+	cords = predictions[(prob_size + conf_size) : ]
+	probs = probs.reshape([SS, C])
+	confs = confs.reshape([SS, B])
+	cords = cords.reshape([SS, B, 4])
+
+	for grid in range(SS):
+		for b in range(B):
+			new_box   = BoundBox(C)
+			new_box.c =  confs[grid, b]
+			new_box.x = (cords[grid, b, 0] + grid %  S) / S
+			new_box.y = (cords[grid, b, 1] + grid // S) / S
+			new_box.w =  cords[grid, b, 2] ** 2
+			new_box.h =  cords[grid, b, 3] ** 2
+			new_box.id = '{}-{}'.format(grid, b)
+			for c in range(C):
+				new_box.probs[c] = new_box.c * probs[grid, c]
+			boxes.append(new_box)
+
+	# non max suppress boxes
+	if True:
+		for c in range(C):
+			for i in range(len(boxes)): boxes[i].class_num = c
+			boxes = sorted(boxes, cmp=prob_compare)
+			for i in range(len(boxes)):
+				boxi = boxes[i]
+				if boxi.probs[c] == 0: continue
+				for j in range(i + 1, len(boxes)):
+					boxj = boxes[j]
+					boxij = box_intersection(boxi, boxj)
+					boxja = boxj.w * boxj.h
+					apart = boxij / boxja
+					if apart >= .5:
+						if boxi.probs[c] > boxj.probs[c]:
+							boxes[j].probs[c] = 0.
+						else:
+							boxes[i].probs[c] = 0.
+
+	imgcv = cv2.imread(img_path)
+	if flip: imgcv = cv2.flip(imgcv, 1)
+	print img_path
+	h, w, _ = imgcv.shape
+	for b in boxes:
+		max_indx = np.argmax(b.probs)
+		max_prob = b.probs[max_indx]
+		label = 'object' * int(C < 2)
+		label += labels[max_indx] * int(C > 1)
+		if (max_prob > threshold):
+			left  = int ((b.x - b.w/2.) * w)
+			right = int ((b.x + b.w/2.) * w)
+			top   = int ((b.y - b.h/2.) * h)
+			bot   = int ((b.y + b.h/2.) * h)
+			if left  < 0    :  left = 0
+			if right > w - 1: right = w - 1
+			if top   < 0    :   top = 0
+			if bot   > h - 1:   bot = h - 1
+			cv2.rectangle(imgcv, 
+				(left, top), (right, bot), 
+				colors[max_indx], 2)
+			mess = '{}:{:.3f}'.format(label, max_prob)
+			cv2.putText(imgcv, mess, (left, top - 10), 
+				0, 1e-3 * h, colors[max_indx])
+
+	img_name = 'results/{}'.format(
+		img_path.split('/')[-1].split('.')[0])
+	cv2.imwrite(img_name + flip * '_' + '.jpg', imgcv)