Skip to content

Instantly share code, notes, and snippets.

What would you like to do?
Caffe feature extractor
import numpy as np
import os, sys, getopt
# Main path to your caffe installation
caffe_root = '/path/to/your/caffe/'
# Model prototxt file
model_prototxt = caffe_root + 'models/bvlc_googlenet/deploy.prototxt'
# Model caffemodel file
model_trained = caffe_root + 'models/bvlc_googlenet/bvlc_googlenet.caffemodel'
# File containing the class labels
imagenet_labels = caffe_root + 'data/ilsvrc12/synset_words.txt'
# Path to the mean image (used for input processing)
mean_path = caffe_root + 'python/caffe/imagenet/ilsvrc_2012_mean.npy'
# Name of the layer we want to extract
layer_name = 'pool5/7x7_s1'
sys.path.insert(0, caffe_root + 'python')
import caffe
def main(argv):
inputfile = ''
outputfile = ''
opts, args = getopt.getopt(argv,"hi:o:",["ifile=","ofile="])
except getopt.GetoptError:
print ' -i <inputfile> -o <outputfile>'
for opt, arg in opts:
if opt == '-h':
print ' -i <inputfile> -o <outputfile>'
elif opt in ("-i"):
inputfile = arg
elif opt in ("-o"):
outputfile = arg
print 'Reading images from "', inputfile
print 'Writing vectors to "', outputfile
# Setting this to CPU, but feel free to use GPU if you have CUDA installed
# Loading the Caffe model, setting preprocessing parameters
net = caffe.Classifier(model_prototxt, model_trained,
image_dims=(256, 256))
# Loading class labels
with open(imagenet_labels) as f:
labels = f.readlines()
# This prints information about the network layers (names and sizes)
# You can uncomment this, to have a look inside the network and choose which layer to print
#print [(k, for k, v in net.blobs.items()]
# Processing one image at a time, printint predictions and writing the vector to a file
with open(inputfile, 'r') as reader:
with open(outputfile, 'w') as writer:
for image_path in reader:
image_path = image_path.strip()
input_image =
prediction = net.predict([input_image], oversample=False)
print os.path.basename(image_path), ' : ' , labels[prediction[0].argmax()].strip() , ' (', prediction[0][prediction[0].argmax()] , ')'
np.savetxt(writer, net.blobs[layer_name].data[0].reshape(1,-1), fmt='%.8g')
if __name__ == "__main__":
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.