classifier.py 文件源码

def load_model(model="bvlc_reference_caffenet"):
    # Path to caffe install
    # caffe_root = '~/deep-learning/caffe'
    caffe_root = "~/git/caffe"
    caffe_root = os.path.expanduser(caffe_root)

    # Set the right paths to your model definition file, pretrained model weights
    # and labels file. This example uses the pre-trained ILSVRC12 image classifier
    # CaffeNet model.
    # You can download it by following the installation instructions steps under
    # http://caffe.berkeleyvision.org/model_zoo.htmli
    MODEL_FILE = caffe_root + ('/models/%s/deploy.prototxt' % model)
    PRETRAINED = caffe_root + ('/models/%s/%s.caffemodel' % (model, model))
    LABELS_FILE = caffe_root + '/data/ilsvrc12/synset_words.txt'
    MEAN_FILE = caffe_root + "/python/caffe/imagenet/ilsvrc_2012_mean.npy"

    # load the network via the cafe.Classifier() method
    net = caffe.Classifier(MODEL_FILE, PRETRAINED,
                           mean=np.load(MEAN_FILE).mean(1).mean(1),
                           channel_swap=(2, 1, 0),
                           raw_scale=255,
                           image_dims=(256, 256))

    # get labels from according file
    labels = []
    with open(LABELS_FILE) as f:
        labels = pd.DataFrame([
            {
                'synset_id': l.strip().split(' ')[0],
                'name': ' '.join(l.strip().split(' ')[1:]).split(',')[0]
            }
            for l in f.readlines()])
        labels = labels.sort_values('synset_id')['name'].values

    return net, labels

# Worker function that takes a frame to be classified from an input queue and
# returns the classification result