!wget https://www.dropbox.com/s/blrajqirr1p31v0/cifar10_nin.caffemodel 
!wget https://gist.githubusercontent.com/ebenolson/91e2cfa51fdb58782c26/raw/b015b7403d87b21c6d2e00b7ec4c0880bbeb1f7e/model.prototxt

"""
Explanation: Introduction
This example demonstrates how to convert a network from Caffe's Model Zoo for use with Lasagne. We will be using the NIN model trained for CIFAR10.
We will create a set of Lasagne layers corresponding to the Caffe model specification (prototxt), then copy the parameters from the caffemodel file into our model.
Final product
If you just want to try the final result, you can download the pickled weights here
Converting from Caffe to Lasagne
Download the required files
First we download cifar10_nin.caffemodel and model.prototxt. The supplied train_val.prototxt was modified to replace the data layers with an input specification, and remove the unneeded loss/accuracy layers.
End of explanation
"""


import caffe

"""
Explanation: Import Caffe
To load the saved parameters, we'll need to have Caffe's Python bindings installed.
End of explanation
"""


net_caffe = caffe.Net('model.prototxt', 'cifar10_nin.caffemodel', caffe.TEST)

"""
Explanation: Load the pretrained Caffe network
End of explanation
"""


import lasagne
from lasagne.layers import InputLayer, DropoutLayer, FlattenLayer
from lasagne.layers.dnn import Conv2DDNNLayer as ConvLayer
from lasagne.layers import Pool2DLayer as PoolLayer
from lasagne.utils import floatX

"""
Explanation: Import Lasagne
End of explanation
"""


net = {}
net['input'] = InputLayer((None, 3, 32, 32))
net['conv1'] = ConvLayer(net['input'], num_filters=192, filter_size=5, pad=2, flip_filters=False)
net['cccp1'] = ConvLayer(net['conv1'], num_filters=160, filter_size=1, flip_filters=False)
net['cccp2'] = ConvLayer(net['cccp1'], num_filters=96, filter_size=1, flip_filters=False)
net['pool1'] = PoolLayer(net['cccp2'], pool_size=3, stride=2, mode='max', ignore_border=False)
net['drop3'] = DropoutLayer(net['pool1'], p=0.5)
net['conv2'] = ConvLayer(net['drop3'], num_filters=192, filter_size=5, pad=2, flip_filters=False)
net['cccp3'] = ConvLayer(net['conv2'], num_filters=192, filter_size=1, flip_filters=False)
net['cccp4'] = ConvLayer(net['cccp3'], num_filters=192, filter_size=1, flip_filters=False)
net['pool2'] = PoolLayer(net['cccp4'], pool_size=3, stride=2, mode='average_exc_pad', ignore_border=False)
net['drop6'] = DropoutLayer(net['pool2'], p=0.5)
net['conv3'] = ConvLayer(net['drop6'], num_filters=192, filter_size=3, pad=1, flip_filters=False)
net['cccp5'] = ConvLayer(net['conv3'], num_filters=192, filter_size=1, flip_filters=False)
net['cccp6'] = ConvLayer(net['cccp5'], num_filters=10, filter_size=1, flip_filters=False)
net['pool3'] = PoolLayer(net['cccp6'], pool_size=8, mode='average_exc_pad', ignore_border=False)
net['output'] = lasagne.layers.FlattenLayer(net['pool3'])

"""
Explanation: Create a Lasagne network
Layer names match those in model.prototxt
End of explanation
"""


layers_caffe = dict(zip(list(net_caffe._layer_names), net_caffe.layers))

for name, layer in net.items():
    try:
        layer.W.set_value(layers_caffe[name].blobs[0].data)
        layer.b.set_value(layers_caffe[name].blobs[1].data)       
    except AttributeError:
        continue

"""
Explanation: Copy the parameters from Caffe to Lasagne
End of explanation
"""


import numpy as np
import pickle
import matplotlib.pyplot as plt

%matplotlib inline

"""
Explanation: Trying it out
Let's see if that worked. 
Import numpy and set up plotting
End of explanation
"""


!wget https://s3.amazonaws.com/lasagne/recipes/pretrained/cifar10/cifar10.npz

data = np.load('cifar10.npz')

"""
Explanation: Download some test data
Since the network expects ZCA whitened and normalized input, we'll download a preprocessed portion (1000 examples) of the CIFAR10 test set.
End of explanation
"""


prob = np.array(lasagne.layers.get_output(net['output'], floatX(data['whitened']), deterministic=True).eval())
predicted = np.argmax(prob, 1)

"""
Explanation: Make predictions on the test data
End of explanation
"""


accuracy = np.mean(predicted == data['labels'])
print(accuracy)

"""
Explanation: Check our accuracy
We expect around 90%
End of explanation
"""


net_caffe.blobs['data'].reshape(1000, 3, 32, 32)
net_caffe.blobs['data'].data[:] = data['whitened']
prob_caffe = net_caffe.forward()['pool3'][:,:,0,0]

np.allclose(prob, prob_caffe)

"""
Explanation: Double check
Let's compare predictions against Caffe
End of explanation
"""


def make_image(X):
    im = np.swapaxes(X.T, 0, 1)
    im = im - im.min()
    im = im * 1.0 / im.max()
    return im

plt.figure(figsize=(16, 5))
for i in range(0, 10):
    plt.subplot(1, 10, i+1)
    plt.imshow(make_image(data['raw'][i]), interpolation='nearest')
    true = data['CLASSES'][data['labels'][i]]
    pred = data['CLASSES'][predicted[i]]
    color = 'green' if true == pred else 'red'
    plt.text(0, 0, true, color='black', bbox=dict(facecolor='white', alpha=1))
    plt.text(0, 32, pred, color=color, bbox=dict(facecolor='white', alpha=1))

    plt.axis('off')

"""
Explanation: Graph some images and predictions
End of explanation
"""


import pickle

values = lasagne.layers.get_all_param_values(net['output'])
pickle.dump(values, open('model.pkl', 'w'))

"""
Explanation: Save our model
Let's save the weights in pickle format, so we don't need Caffe next time
End of explanation
"""


