finish search engine

server/build_index.py

@@ -5,15 +5,13 @@
 
 data_dir = '../data/'
 
-mData = MetaData()
+mData = MetaData('image_db')
 for each in os.walk(data_dir):
 	dirname, dirs, files = each
 	if not files:
 		continue
 	for f in files:
 		if f.endswith('.jpg') and not f.endswith('_contour.jpg'):
 			mData.add(os.path.join(dirname, f))
-
-with open('db.pkl', 'wb') as f:
-	pickle.dump(mData, f)
+mData.save()
 print 'database saved'

server/metadata.py

@@ -1,16 +1,44 @@
 from numpy import *
 import scipy.io as spio
 import os
+try:
+	import cPickle as pickle
+except:
+	import pickle
 
 from pfutils import *
 
 class MetaData:
-	def __init__(self):
+	def __init__(self, dbname = 'image_db'):
+		self.dbname = dbname
+
 		self.tot = 0
 		self.i2name = {}
 		self.name2i = {}
 		self.invIdx = {}
 		self.i2olen = {}
+		self.i2hmap = {}
+
+	def load(self, dbname = 'image_db'):
+		self.dbname = dbname
+		with open(dbname + '.pkl', 'rb') as f:
+			self.tot = pickle.load(f)
+			self.i2name = pickle.load(f)
+			self.name2i = pickle.load(f)
+			self.invIdx = pickle.load(f)
+			self.i2olen = pickle.load(f)
+		with open(dbname + '.hmap.pkl', 'rb') as f:
+			self.i2hmap = pickle.load(f)
+
+	def save(self):
+		with open(self.dbname + '.pkl', 'wb') as f:
+			pickle.dump(self.tot, f, protocol = pickle.HIGHEST_PROTOCOL)
+			pickle.dump(self.i2name, f, protocol = pickle.HIGHEST_PROTOCOL)
+			pickle.dump(self.name2i, f, protocol = pickle.HIGHEST_PROTOCOL)
+			pickle.dump(self.invIdx, f, protocol = pickle.HIGHEST_PROTOCOL)
+			pickle.dump(self.i2olen, f, protocol = pickle.HIGHEST_PROTOCOL)
+		with open(self.dbname + '.hmap.pkl', 'wb') as f:
+			pickle.dump(self.i2hmap, f, protocol = pickle.HIGHEST_PROTOCOL)
 
 	def add(self, name):
 		# check name is '*.jpg', not contour image (not '*_contour.jpg') and was preprocessed (has '*.mat')
@@ -29,6 +57,10 @@ def add(self, name):
 		data = spio.loadmat(matname)
 		pb = data['pb']
 
+		pb = im2bw(pb).astype(int8)
+		unused, hmap = hitMap(pb)
+		self.i2hmap[ind] = hmap
+
 		ocm = extractOCM(pb)
 		self.i2olen[ind] = len(ocm)
 		for each in ocm:

server/pfutils.py

@@ -35,7 +35,6 @@ def pyangle2c(x):
 angle2c = frompyfunc(pyangle2c, 1, 1)
 
 def extractOCM(pb):
-	pb = im2bw(pb).astype(int8)
 	px, py = gradient(pb)
 	pa = arctan2(py, px)
 	pa = angle2c(pa)
@@ -46,3 +45,39 @@ def extractOCM(pb):
 				ocm.append((i, j, pa[i][j]))
 	return ocm
 
+# using the same radius (8) in CAO Yang's master thesis
+radius = 8
+
+__dx = [ 1,-1, 0, 0]
+__dy = [ 0, 0, 1,-1]
+
+def hitMap(pb):
+	q = {}
+	ocm = extractOCM(pb)
+
+	hmap = {}
+	for theta in range(1, 7):
+		b = zeros(pb.shape, dtype = int8)
+		q[theta] = []
+
+		for x, y, t in ocm:
+			if t == theta:
+				q[theta].append((x, y, 0))
+				b[x, y] = 1
+
+		i = 0
+		while i < len(q[theta]):
+			tx, ty, td = q[theta][i]
+			if td < radius:
+				for k in range(4):
+					x = tx + __dx[k]
+					y = ty + __dy[k]
+					if 0 <= x < pb.shape[0] and 0 <= y < pb.shape[1] and not b[x][y]:
+						b[x][y] = 1
+						q[theta].append((x, y, td + 1))
+			i += 1
+
+		hmap[theta] = b
+
+	return q, hmap
+

server/query.py

@@ -1,63 +1,17 @@
-import pickle
 import scipy as sp
 import operator
 from pylab import *
 from numpy import *
 from PIL import Image
+from time import clock
 
 from metadata import *
 from pfutils import *
 
-debugging = False
-
-# using the same radius in CAO Yang's master thesis
-radius = 3
-
-__dx = [ 1,-1, 0, 0]
-__dy = [ 0, 0, 1,-1]
-
-def hitMap(pb):
-	q = {}
-	ocm = extractOCM(pb)
-	row = range(pb.shape[0])
-	col = range(pb.shape[1])
-	debug(str(pb.shape[0] * pb.shape[1]))
-
-	for theta in range(1, 7):
-		b = zeros(pb.shape)
-		q[theta] = []
-
-		for x, y, t in ocm:
-			if t == theta:
-				q[theta].append((x, y, 0))
-				b[x, y] = 1
-
-		debug('init length: %d' % len(q[theta]))
-		i = 0
-		while i < len(q[theta]):
-			tx, ty, td = q[theta][i]
-			if td < radius:
-				for k in range(4):
-					x = tx + __dx[k]
-					y = ty + __dy[k]
-					if x in row and y in col and not b[x][y]:
-						b[x][y] = 1
-						q[theta].append((x, y, td + 1))
-			i += 1
-		debug('length: %d' % i)
-	if debugging:
-		figure()
-		gray()
-		imshow(b)
-		figure()
-		gray()
-		imshow(pb)
-		show()
-	return q
-
 def getMatch(pb):
 	debug('calculating hit map')
-	hm = hitMap(pb)
+	pb = im2bw(pb).astype(int8)
+	hm, unused = hitMap(pb)
 
 	debug('searching')
 	match = {}
@@ -78,35 +32,63 @@ def getMatch(pb):
 
 	return match, sorted_m
 
+def debug_t(msg):
+	debug('time used: %f' % (clock() - debug_t.tstamp))
+	debug(msg)
+	debug_t.tstamp = clock()
+debug_t.tstamp = clock()
+
 def test():
 	debug('loading database')
-	#with open('db.pkl', 'rb') as f:
-	#	mData = pickle.load(f)
+	# mData = MetaData()
+	# mData.load()
+
+	debug_t('loading mat file')
 
-	debug('loading mat file')
-	matname = '../data/swan/5.mat'
+	matname = '../data/teapot/22.mat'
 	data = sp.io.loadmat(matname)
 	pb = data['pb']
 
+	debug_t('one side match')
+
 	match, sorted_m = getMatch(pb)
-	debug('length of sorted_m: %d' % len(sorted_m))
-	for each in sorted_m(:max(len, 1000)):
-		pass
 
-	if debugging:
+	debug_t('finish one side match')
+
+	if True:
 		name = matname[:-3] + 'jpg'
-		debug(str(sorted_m[:10]))
 		debug('origin: image %s score: %f' % (name, match[mData.name2i[name]]))
 		figure()
 		imshow(array(Image.open(name)))
-		for i in range(10):
-			name = mData.i2name[sorted_m[i][0]]
-			debug('name: %s score: %f' % (name, sorted_m[i][1]))
-			#figure()
-			#imshow(array(Image.open(name)))
-		show()
+		axis('off')
+
+	debug_t('two side match')
+
+	pb = im2bw(pb).astype(int8)
+	qocm = extractOCM(pb)
+	for ind, score1 in sorted_m[:min(len(sorted_m), 2000)]:
+		hmap = mData.i2hmap[ind]
+		score2 = 0
+		for x, y, theta in qocm:
+			if x < hmap[theta].shape[0] and y < hmap[theta].shape[1] and hmap[theta][x][y]:
+				score2 += 1
+		score2 = float(score2) / len(qocm)
+		match[ind] = sqrt(score1 * score2)
 
-	return
+	sorted_m = sorted(match.iteritems(), key = operator.itemgetter(1), reverse = True)
+
+	debug_t('finish two side match')
+
+	if True:
+		figure()
+		lenRes = 30
+		for i in range(lenRes):
+			imname = mData.i2name[sorted_m[i][0]]
+			debug('name: %s score: %f' % (imname, sorted_m[i][1]))
+			subplot(5, 6, i + 1)
+			imshow(array(Image.open(imname)))
+			axis('off')
+		show()
 
 if __name__ == '__main__':
 	test()