Skip to content

Instantly share code, notes, and snippets.

@kgrm
Created July 10, 2016 17:49
Show Gist options
  • Save kgrm/8902e13cdea64127c5d15104d8f5acbe to your computer and use it in GitHub Desktop.
Save kgrm/8902e13cdea64127c5d15104d8f5acbe to your computer and use it in GitHub Desktop.
import sys, os, cv2
import caffe
import numpy as np
model_file = "YOURNET.prototxt"
pretrained = "YOURNET.caffemodel"
# Load a pretrained CAFFE neural network and an image,
# use the network to extract features from the image,
# then save the feature vector to a text file.
if len(sys.argv) != 4:
print "ERROR: missing input arguments."
print "Usage: %s mode input output" % sys.argv[0]
print "(see README.txt)"
sys.exit(1)
else:
if "cpu" in sys.argv[1]:
caffe.set_mode_cpu()
net = caffe.Net(model_file, pretrained, caffe.TEST)
elif "gpu" in sys.argv[1]:
caffe.set_mode_gpu()
net = caffe.Net(model_file, pretrained, caffe.TEST)
else:
print "Mode must be cpu or gpu."
sys.exit(1)
img = cv2.imread(sys.argv[2]).transpose((2,0,1)) # convert BGR to RGB
# net.blobs["data"] and net.blobs["fc7"] are, respectively, the inputs
# of the networks you need to provide, and the output of the network
# you're interested in. These are layer names, and will probably differ
# from network to network.
net.blobs["data"].data[...] = img
outs = net.forward(["fc7"])
output_feats = outs["fc7"].reshape((4096,))
with open(sys.argv[3], "w") as f:
for feat in output_feats:
f.write("%1.4f " % feat)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment