pycaffe_example.py

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)