#!/usr/bin/env python2

import os
import pickle

import caffe
caffe.set_mode_cpu()
caffe_root = os.path.normpath(os.path.dirname('%s/../../../../' % caffe.__file__))

import numpy as np

import util

def main():
    extractor = FeatureExtractor()
    cache_features(extractor, '../datasets/Wookie')  # Train set
    cache_features(extractor, '../datasets/Michael') # Test set

def cache_features(extractor, path):
    files = util.filesWithRe(path, r'.*\.(jpg|jpeg|png)$')
    for i, fpath in enumerate(files):
        feats = extractor.get_features(fpath)

        with open('%s.pickle' % fpath, 'w') as f:
            pickle.dump(feats, f)

        print('[%d/%d] Stored features for %s' % (i, len(files), fpath))

class FeatureExtractor:
    def __init__(self):
        caffe.set_mode_cpu()
        net = caffe.Net(caffe_root + '/models/bvlc_googlenet/deploy.prototxt',
                        caffe_root + '/models/bvlc_googlenet/bvlc_googlenet.caffemodel',
                        caffe.TEST)

        # Input preprocessing: 'data' is the name of the input blob == net.inputs[0]
        transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
        transformer.set_transpose('data', (2, 0, 1))

        # Mean pixel
        transformer.set_mean('data', np.load(caffe_root + '/python/caffe/imagenet/ilsvrc_2012_mean.npy').mean(1).mean(1))

        # The reference model operates on images in [0,255] range instead of [0,1]
        transformer.set_raw_scale('data', 255)

        # The reference model has channels in BGR order instead of RGB
        transformer.set_channel_swap('data', (2, 1, 0))

        self.net = net
        self.transformer = transformer

    def get_features(self, img_path):
        # Reshape net for single image input
        self.net.blobs['data'].reshape(1, 3, 224, 224)

        img = self.transformer.preprocess('data', caffe.io.load_image(img_path))
        self.net.blobs['data'].data[...] = img

        out = self.net.forward()

        return {
            'classes': self.net.blobs['loss3/classifier'].data.flatten(),
        }

if __name__ == '__main__':
    main()


@@ -834,29 +834,38 @@
 models/bvlc_
-googl
+reference_caff
 enet/deploy.
@@ -930,28 +930,46 @@
 vlc_
-googlenet/bvlc_googl
+reference_caffenet/bvlc_reference_caff
 enet
@@ -1841,14 +1841,14 @@
 , 22
-4
+7
 , 22
-4
+7
 )%0A%0A 
@@ -2069,24 +2069,11 @@
 bs%5B'
-loss3/classifier
+fc6
 '%5D.d

Use reference ImageNet model from Caffee