This script is for computing the performance of the network on our sketchy benchmark.
In [1]:
import numpy as np
from pylab import *
%matplotlib inline
import os
import sys
First, we need to import caffe. You'll need to have caffe installed, as well as python interface for caffe.
In [2]:
#TODO: specify your caffe root folder here
caffe_root = "X:\caffe_siggraph/caffe-windows-master"
sys.path.insert(0, caffe_root+'/python')
import caffe
Now we can load up the network. You can change the path to your own network here. Make sure to use the matching deploy prototxt files and change the target layer to your layer name.
In [35]:
#TODO: change to your own network and deploying file
PRETRAINED_FILE = '../models/triplet_googlenet/triplet_googlenet_finegrain_final.caffemodel'
sketch_model = '../models/triplet_googlenet/googlenet_sketchdeploy.prototxt'
image_model = '../models/triplet_googlenet/googlenet_imagedeploy.prototxt'
In [36]:
caffe.set_mode_gpu()
#caffe.set_mode_cpu()
sketch_net = caffe.Net(sketch_model, PRETRAINED_FILE, caffe.TEST)
img_net = caffe.Net(image_model, PRETRAINED_FILE, caffe.TEST)
sketch_net.blobs.keys()
Out[36]:
In [38]:
#TODO: set output layer name. You can use sketch_net.blobs.keys() to list all layer
output_layer_sketch = 'pool5/7x7_s1_s'
output_layer_image = 'pool5/7x7_s1_p'
In [12]:
#set the transformer
transformer = caffe.io.Transformer({'data': np.shape(sketch_net.blobs['data'].data)})
transformer.set_mean('data', np.array([104, 117, 123]))
transformer.set_transpose('data',(2,0,1))
transformer.set_channel_swap('data', (2,1,0))
transformer.set_raw_scale('data', 255.0)
In [24]:
#photo paths
photo_paths = 'C:\Users\Patsorn\Documents/notebook_backup/SBIR/photos/'
sketch_paths = 'C:\Users\Patsorn\Documents/notebook_backup/SBIR/sketches/'
In [15]:
#load up test images
with open('../list/test_img_list.txt','r') as my_file:
test_img_list = [c.rstrip() for c in my_file.readlines()]
In [39]:
#extract feature for all test images
feats = []
N = np.shape(test_img_list)[0]
for i,path in enumerate(test_img_list):
imgname = path.split('/')[-1]
imgname = imgname.split('.jpg')[0]
imgcat = path.split('/')[0]
print '\r',str(i+1)+'/'+str(N)+ ' '+'Extracting ' +path+'...',
full_path = photo_paths + path
img = (transformer.preprocess('data', caffe.io.load_image(full_path.rstrip())))
img_in = np.reshape([img],np.shape(sketch_net.blobs['data'].data))
out_img = img_net.forward(data=img_in)
out_img = np.copy(out_img[output_layer_image])
feats.append(out_img)
print 'done',
np.shape(feats)
feats = np.resize(feats,[np.shape(feats)[0],np.shape(feats)[2]]) #quick fixed for size
In [44]:
#build nn pool
from sklearn.neighbors import NearestNeighbors,LSHForest
nbrs = NearestNeighbors(n_neighbors=np.size(feats,0), algorithm='brute',metric='cosine').fit(feats)
In [45]:
#compute score
num_query = 0
count_recall = [0]*1250
sum_rank = 0
sum_class_rank = [0]*125
count_recall_class = np.zeros((125,1250),np.float)
i_coco =-1
for i,img in enumerate(test_img_list):
imgname = img.split('/')[-1]
imgname = imgname.split('.jpg')[0]
imgcat = img.split('/')[0]
sketch_list = os.listdir(sketch_paths+imgcat)
sketch_img_list = [skg for skg in sketch_list if skg.startswith(imgname+'-') and skg.endswith('-5.png')]#change this skg.endswith('-1.png') to the variation you want
for sketch in sketch_img_list:
sketch_path = sketch_paths + imgcat+'/' + sketch
sketch_in = (transformer.preprocess('data', plt.imread(sketch_path)))
sketch_in = np.reshape([sketch_in],np.shape(sketch_net.blobs['data'].data))
query = sketch_net.forward(data=sketch_in)
query=np.copy(query[output_layer_sketch])
distances, indices = nbrs.kneighbors(np.reshape(query,[np.shape(query)[1]]))
num_query = num_query+1
print '\r','...'+sketch+'...',
for j,indice in enumerate(indices[0]):
if indice==i:
#this j is the right one.
count_recall[j] = count_recall[j]+1
print '\r','ranking: '+imgcat+ ' '+sketch + ' found at ' +str(j),
break
cum_count = [0]*1250
sumc = 0
for i,c in enumerate(count_recall):
sumc = sumc + c
cum_count[i] = sumc
print '\nRecall @K=1 = ', 1.00*cum_count[0]/cum_count[-1]