Skip to content

Instantly share code, notes, and snippets.

@volkancirik
Created June 8, 2016 17:33
Show Gist options
  • Save volkancirik/8780878a1a2cc6b1101658b7164b476c to your computer and use it in GitHub Desktop.
Save volkancirik/8780878a1a2cc6b1101658b7164b476c to your computer and use it in GitHub Desktop.
extract cnn features
# Make sure that caffe is on the python path:
caffe_root = '/usr0/bin/caffe/' # this file is expected to be in {caffe_root}/examples
import sys
sys.path.insert(0, caffe_root + 'python')
import caffe
from images import crop_image
import numpy as np
mean = 'caffedata/ilsvrc_2012_mean.npy'
class CNN(object):
def __init__(self, deploy, model, mean=mean, batch_size=100, width=224, height=224):
self.deploy = deploy
self.model = model
self.mean = mean
self.batch_size = batch_size
self.net, self.transformer = self.get_net()
self.net.blobs['data'].reshape(self.batch_size, 3, height, width)
self.width = width
self.height = height
def get_net(self):
caffe.set_mode_gpu()
net = caffe.Net(self.deploy, self.model, caffe.TEST)
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_transpose('data', (2,0,1))
# transformer.set_mean('data', np.load(self.mean).mean(1).mean(1))
transformer.set_mean('data', np.array([103.939, 116.779, 123.68]))
transformer.set_raw_scale('data', 255)
transformer.set_channel_swap('data', (2,1,0))
return net, transformer
def get_features(self, image_list, layers='fc7', layer_sizes=[4096]):
iter_until = len(image_list) + self.batch_size
all_feats = np.zeros([len(image_list)] + layer_sizes)
batch_count = 0
for start, end in zip(range(0, iter_until, self.batch_size), \
range(self.batch_size, iter_until, self.batch_size)):
image_batch_file = image_list[start:end]
image_batch = np.array(map(lambda x: crop_image(x, target_width=self.width, target_height=self.height), image_batch_file))
caffe_in = np.zeros(np.array(image_batch.shape)[[0,3,1,2]], dtype=np.float32)
for idx, in_ in enumerate(image_batch):
caffe_in[idx] = self.transformer.preprocess('data', in_)
out = self.net.forward_all(blobs=[layers], **{'data':caffe_in})
feats = out[layers]
all_feats[start:end] = feats
batch_count += 1
print("processing batch #: %d" % batch_count)
return all_feats
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment