def gpu_task(prototxt, caffemodel, layer, path_images, out, gpu=0):
try:
num_images = len(path_images)
h5f = h5py.File(out, 'w')
# set gpu card
caffe.set_device(gpu)
caffe.set_mode_gpu()
# init NN
net = caffe.Net(prototxt, caffemodel, caffe.TEST)
net.forward()
features = []
image_names = []
for i, path in enumerate(path_images):
print "%d(%d), %s"%((i+1), num_images, os.path.basename(path))
d = opencv_format_img_for_vgg(path, True)
feat = extract_fc_features(net, layer, d)
#features.append(feat.tolist())
features.append(np.array(feat))
image_names.append(os.path.basename(path))
features = np.array(features)
h5f['feats'] = features
h5f['names'] = image_names
h5f.close()
print "gpu %d task has finished..." % (int(gpu))
except:
import traceback
traceback.print_exc()
raise
def query_images(groundtruth_dir, image_dir, dataset, cropped=True):
imgs = []
query_names = []
fake_query_names = []
feats_crop = []
modelDir = "./fc_retrieval/model"
MODEL = "nueral.caffemodel"
PROTO = "deploy.prototxt"
caffemodel = os.path.join(modelDir, MODEL)
prototxt = os.path.join(modelDir, PROTO)
# set gpu card
layer = 'fc6'
caffe.set_device(7)
caffe.set_mode_gpu()
# init NN
net = caffe.Net(prototxt, caffemodel, caffe.TEST)
net.forward()
for f in glob.iglob(os.path.join(groundtruth_dir, '*_query.txt')):
fake_query_name = os.path.splitext(os.path.basename(f))[0].replace(
'_query', '')
fake_query_names.append(fake_query_name)
query_name, x, y, w, h = open(f).read().strip().split(' ')
if dataset == 'oxford':
query_name = query_name.replace('oxc1_', '')
query_names.append('%s.jpg' % query_name)
img = cv2.imread(os.path.join(image_dir, '%s.jpg' % query_name),
1) # BGR
if cropped:
x, y, w, h = map(float, (x, y, w, h))
x, y, w, h = map(lambda d: int(round(d)), (x, y, w, h))
else:
x, y, w, h = (0, 0, img.shape[1], img.shape[0])
img = img[y:y + h, x:x + w]
img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR)
d = np.float32(img)
# VggNet
d -= np.array((104.00698793, 116.66876762, 122.67891434))
d = d.transpose((2, 0, 1))
feat = extract_fc_features(net, layer, d)
# L2-normalize feature
feat = normalize(feat, copy=False)
#feats_crop.append(feat.tolist())
feats_crop.append(np.array(feat))
imgs.append(img)
return imgs, feats_crop, query_names, fake_query_names