import caffe


"""
Explanation: Introduction
This example demonstraites how to convert Caffe pretrained ResNet-50 model from https://github.com/KaimingHe/deep-residual-networks (firstly described in http://arxiv.org/pdf/1512.03385v1.pdf) into Theano/Lasagne format.
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 (like <a href="https://github.com/Lasagne/Recipes/blob/master/examples/Using%20a%20Caffe%20Pretrained%20Network%20-%20CIFAR10.ipynb">here</a>).
This notebook produce resnet50.pkl file, which contains dictionary with following foelds:
 * values: numpy array with parameters of the model
 * synset_words: labels of classes
 * mean_image: mean image which should be subtracted from each input image
This file can be used for initialization of weights of the model created by modelzoo/resnet50.py.
License
Same as in parent project https://github.com/KaimingHe/deep-residual-networks/blob/master/LICENSE
Requirements
Download the required files
<a href="https://onedrive.live.com/?authkey=%21AAFW2-FVoxeVRck&id=4006CBB8476FF777%2117887&cid=4006CBB8476FF777">Here</a> you can find folder with caffe/proto files, we need followings to be stored in ./:
 * ResNet-50-deploy.prototxt contains architecture of ResNet-50 in proto format
 * ResNet-50-model.caffemodel is proto serialization of model parameters
 * ResNet_mean.binaryproto contains mean values
Imports
We need caffe to load weights and compare results
End of explanation
"""


import lasagne
from lasagne.utils import floatX
from lasagne.layers import InputLayer
from lasagne.layers import Conv2DLayer as ConvLayer # can be replaced with dnn layers
from lasagne.layers import BatchNormLayer
from lasagne.layers import Pool2DLayer as PoolLayer
from lasagne.layers import NonlinearityLayer
from lasagne.layers import ElemwiseSumLayer
from lasagne.layers import DenseLayer
from lasagne.nonlinearities import rectify, softmax

"""
Explanation: We need a lot of building blocks from Lasagne to build network
End of explanation
"""


%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['figure.figsize'] = 8, 6
import io
import urllib
import skimage.transform
from IPython.display import Image
import pickle

"""
Explanation: Helper modules, some of them will help us to download images and plot them
End of explanation
"""


Image(filename='images/head.png', width='40%')

"""
Explanation: Build Lasagne model
BatchNormalization issue in caffe
Caffe doesn't have correct BN layer as described in https://arxiv.org/pdf/1502.03167.pdf:
 * it can collect datasets mean ($\hat{\mu}$) and variance ($\hat{\sigma}^2$)
 * it can't fit $\gamma$ and $\beta$ parameters to scale and shift standardized distribution of feature in following formula: $\hat{x}_i = \dfrac{x_i - \hat{\mu}_i}{\sqrt{\hat{\sigma}_i^2 + \epsilon}}\cdot\gamma + \beta$
To fix this issue, <a href="https://github.com/KaimingHe/deep-residual-networks">here</a> authors use such BN layer followed by Scale layer, which can fit scale and shift parameters, but can't standardize data:
<pre>
layer {
    bottom: "res2a_branch1"
    top: "res2a_branch1"
    name: "bn2a_branch1"
    type: "BatchNorm"
    batch_norm_param {
        use_global_stats: true
    }
}

layer {
    bottom: "res2a_branch1"
    top: "res2a_branch1"
    name: "scale2a_branch1"
    type: "Scale"
    scale_param {
        bias_term: true
    }
}
</pre>

In Lasagne we have correct BN layer, so we do not need use such a trick.
Replicated blocks
Simple blocks
ResNet contains a lot of similar replicated blocks, lets call them simple blocks, which have one of two architectures:
 * Convolution $\rightarrow$ BN $\rightarrow$ Nonlinearity
 * Convolution $\rightarrow$ BN
http://ethereon.github.io/netscope/#/gist/2f702ea9e05900300462102a33caff9c
End of explanation
"""


Image(filename='images/conv1x1.png', width='40%')

"""
Explanation: We can increase, decrease or keep same dimensionality of data using such blocks. In ResNet-50 only several transformation are used.
Keep shape with 1x1 convolution
We can apply nonlinearity transformation from (None, 64, 56, 56) to (None, 64, 56, 56) if we apply simple block with following parameters (look at the origin of a network after first pool layer):
 * num_filters: same as parent has
 * filter_size: 1
 * stride: 1
 * pad: 0
End of explanation
"""


Image(filename='images/conv3x3.png', width='40%')

"""
Explanation: Keep shape with 3x3 convolution
Also we can apply nonlinearity transformation from (None, 64, 56, 56) to (None, 64, 56, 56) if we apply simple block with following parameters (look at the middle of any residual blocks):
 * num_filters: same as parent has
 * filter_size: 3x3
 * stride: 1
 * pad: 1
End of explanation
"""


Image(filename='images/increase_fn.png', width='40%')

"""
Explanation: Increase shape using number of filters
We can nonlinearly increase shape from (None, 64, 56, 56) to (None, 256, 56, 56) if we apply simple block with following parameters (look at the last simple block of any risidual block):
 * num_filters: four times greater then parent has
 * filter_size: 1x1
 * stride: 1
 * pad: 0
End of explanation
"""


Image(filename='images/decrease_fn.png', width='40%')

"""
Explanation: Increase shape using number of filters
We can nonlinearly decrease shape from (None, 256, 56, 56) to (None, 64, 56, 56) if we apply simple block with following parameters (look at the first simple block of any risidual block without left branch):
 * num_filters: four times less then parent has
 * filter_size: 1x1
 * stride: 1
 * pad: 0
End of explanation
"""


Image(filename='images/decrease_fnstride.png', width='40%')

"""
Explanation: Increase shape using number of filters
We can also nonlinearly decrease shape from (None, 256, 56, 56) to (None, 128, 28, 28) if we apply simple block with following parameters (look at the first simple block of any risidual block with left branch):
 * num_filters: two times less then parent has
 * filter_size: 1x1
 * stride: 2
 * pad: 0
End of explanation
"""


def build_simple_block(incoming_layer, names,
                       num_filters, filter_size, stride, pad, 
                       use_bias=False, nonlin=rectify):
    """Creates stacked Lasagne layers ConvLayer -> BN -> (ReLu)
    
    Parameters:
    ----------
    incoming_layer : instance of Lasagne layer
        Parent layer
    
    names : list of string
        Names of the layers in block
        
    num_filters : int
        Number of filters in convolution layer
        
    filter_size : int
        Size of filters in convolution layer
        
    stride : int
        Stride of convolution layer
        
    pad : int
        Padding of convolution layer
        
    use_bias : bool
        Whether to use bias in conlovution layer
        
    nonlin : function
        Nonlinearity type of Nonlinearity layer
        
    Returns
    -------
    tuple: (net, last_layer_name)
        net : dict
            Dictionary with stacked layers
        last_layer_name : string
            Last layer name
    """
    net = []
    net.append((
            names[0], 
            ConvLayer(incoming_layer, num_filters, filter_size, pad, stride, 
                      flip_filters=False, nonlinearity=None) if use_bias 
            else ConvLayer(incoming_layer, num_filters, filter_size, stride, pad, b=None, 
                           flip_filters=False, nonlinearity=None)
        ))
    
    net.append((
            names[1], 
            BatchNormLayer(net[-1][1])
        ))
    if nonlin is not None:
        net.append((
            names[2], 
            NonlinearityLayer(net[-1][1], nonlinearity=nonlin)
        ))
    
    return dict(net), net[-1][0]

"""
Explanation: Following function creates simple block
End of explanation
"""


Image(filename='images/left_branch.png', width='40%')

Image(filename='images/no_left_branch.png', width='40%')

simple_block_name_pattern = ['res%s_branch%i%s', 'bn%s_branch%i%s', 'res%s_branch%i%s_relu']

def build_residual_block(incoming_layer, ratio_n_filter=1.0, ratio_size=1.0, has_left_branch=False, 
                         upscale_factor=4, ix=''):
    """Creates two-branch residual block
    
    Parameters:
    ----------
    incoming_layer : instance of Lasagne layer
        Parent layer
    
    ratio_n_filter : float
        Scale factor of filter bank at the input of residual block
        
    ratio_size : float
        Scale factor of filter size
        
    has_left_branch : bool
        if True, then left branch contains simple block
        
    upscale_factor : float
        Scale factor of filter bank at the output of residual block
        
    ix : int
        Id of residual block
        
    Returns
    -------
    tuple: (net, last_layer_name)
        net : dict
            Dictionary with stacked layers
        last_layer_name : string
            Last layer name
    """
    net = {}
    
    # right branch
    net_tmp, last_layer_name = build_simple_block(
        incoming_layer, map(lambda s: s % (ix, 2, 'a'), simple_block_name_pattern),
        int(lasagne.layers.get_output_shape(incoming_layer)[1]*ratio_n_filter), 1, int(1.0/ratio_size), 0)
    net.update(net_tmp)
    
    net_tmp, last_layer_name = build_simple_block(
        net[last_layer_name], map(lambda s: s % (ix, 2, 'b'), simple_block_name_pattern),
        lasagne.layers.get_output_shape(net[last_layer_name])[1], 3, 1, 1)
    net.update(net_tmp)
    
    net_tmp, last_layer_name = build_simple_block(
        net[last_layer_name], map(lambda s: s % (ix, 2, 'c'), simple_block_name_pattern),
        lasagne.layers.get_output_shape(net[last_layer_name])[1]*upscale_factor, 1, 1, 0,
        nonlin=None)
    net.update(net_tmp)
    
    right_tail = net[last_layer_name]
    left_tail = incoming_layer
    
    # left branch
    if has_left_branch:
        net_tmp, last_layer_name = build_simple_block(
            incoming_layer, map(lambda s: s % (ix, 1, ''), simple_block_name_pattern),
            int(lasagne.layers.get_output_shape(incoming_layer)[1]*4*ratio_n_filter), 1, int(1.0/ratio_size), 0,
            nonlin=None)
        net.update(net_tmp)
        left_tail = net[last_layer_name]
        
    net['res%s' % ix] = ElemwiseSumLayer([left_tail, right_tail], coeffs=1)
    net['res%s_relu' % ix] = NonlinearityLayer(net['res%s' % ix], nonlinearity=rectify)
    
    return net, 'res%s_relu' % ix

"""
Explanation: Residual blocks
ResNet also contains several residual blockes built from simple blocks, each of them have two branches; left branch sometimes contains simple block, sometimes not. Each block ends with Elementwise sum layer followed by ReLu nonlinearity.  
http://ethereon.github.io/netscope/#/gist/410e7e48fa1e5a368ee7bca5eb3bf0ca
End of explanation
"""


net = {}
net['input'] = InputLayer((None, 3, 224, 224))
sub_net, parent_layer_name = build_simple_block(
    net['input'], ['conv1', 'bn_conv1', 'conv1_relu'],
    64, 7, 3, 2, use_bias=True)
net.update(sub_net)
net['pool1'] = PoolLayer(net[parent_layer_name], pool_size=3, stride=2, pad=0, mode='max', ignore_border=False)

"""
Explanation: Gathering everighting together
Create head of the network (everithing before first residual block)
End of explanation
"""


block_size = list('abc')
parent_layer_name = 'pool1'
for c in block_size:
    if c == 'a':
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1, 1, True, 4, ix='2%s' % c)
    else:
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/4, 1, False, 4, ix='2%s' % c)
    net.update(sub_net)
    
block_size = list('abcd')
for c in block_size:
    if c == 'a':
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/2, 1.0/2, True, 4, ix='3%s' % c)
    else:
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/4, 1, False, 4, ix='3%s' % c)
    net.update(sub_net)
    
block_size = list('abcdef')
for c in block_size:
    if c == 'a':
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/2, 1.0/2, True, 4, ix='4%s' % c)
    else:
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/4, 1, False, 4, ix='4%s' % c)
    net.update(sub_net)
    
block_size = list('abc')
for c in block_size:
    if c == 'a':
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/2, 1.0/2, True, 4, ix='5%s' % c)
    else:
        sub_net, parent_layer_name = build_residual_block(net[parent_layer_name], 1.0/4, 1, False, 4, ix='5%s' % c)
    net.update(sub_net)

"""
Explanation: Create four groups of residual blocks
End of explanation
"""


net['pool5'] = PoolLayer(net[parent_layer_name], pool_size=7, stride=1, pad=0, 
                         mode='average_exc_pad', ignore_border=False)
net['fc1000'] = DenseLayer(net['pool5'], num_units=1000, nonlinearity=None)
net['prob'] = NonlinearityLayer(net['fc1000'], nonlinearity=softmax)

print 'Total number of layers:', len(lasagne.layers.get_all_layers(net['prob']))

"""
Explanation: Create tail of the network (everighting after last resudual block)
End of explanation
"""


net_caffe = caffe.Net('./ResNet-50-deploy.prototxt', './ResNet-50-model.caffemodel', caffe.TEST)
layers_caffe = dict(zip(list(net_caffe._layer_names), net_caffe.layers))
print 'Number of layers: %i' % len(layers_caffe.keys())

"""
Explanation: Transfer weights from caffe to lasagne
Load pretrained caffe model
End of explanation
"""


for name, layer in net.items():    
    if name not in layers_caffe:
        print name, type(layer).__name__
        continue
    if isinstance(layer, BatchNormLayer):
        layer_bn_caffe = layers_caffe[name]
        layer_scale_caffe = layers_caffe['scale' + name[2:]]
        layer.gamma.set_value(layer_scale_caffe.blobs[0].data)
        layer.beta.set_value(layer_scale_caffe.blobs[1].data)
        layer.mean.set_value(layer_bn_caffe.blobs[0].data)
        layer.inv_std.set_value(1/np.sqrt(layer_bn_caffe.blobs[1].data) + 1e-4)
        continue
    if isinstance(layer, DenseLayer):
        layer.W.set_value(layers_caffe[name].blobs[0].data.T)
        layer.b.set_value(layers_caffe[name].blobs[1].data)
        continue
    if len(layers_caffe[name].blobs) > 0:
        layer.W.set_value(layers_caffe[name].blobs[0].data)
    if len(layers_caffe[name].blobs) > 1:
        layer.b.set_value(layers_caffe[name].blobs[1].data)

"""
Explanation: Copy weights
There is one more issue with BN layer: caffa stores variance $\sigma^2$, but lasagne stores inverted standard deviation $\dfrac{1}{\sigma}$, so we need make simple transfommation to handle it.
Other issue reffers to weights ofthe dense layer, in caffe it is transposed, we should handle it too.
End of explanation
"""


with open('./imagenet_classes.txt', 'r') as f:
    classes = map(lambda s: s.strip(), f.readlines())

"""
Explanation: Testing
Read ImageNet synset
End of explanation
"""


index = urllib.urlopen('http://www.image-net.org/challenges/LSVRC/2012/ori_urls/indexval.html').read()
image_urls = index.split('<br>')
np.random.seed(23)
np.random.shuffle(image_urls)
image_urls = image_urls[:100]

"""
Explanation: Download some image urls for recognition
End of explanation
"""


blob = caffe.proto.caffe_pb2.BlobProto()
data = open('./ResNet_mean.binaryproto', 'rb').read()
blob.ParseFromString(data)
mean_values = np.array(caffe.io.blobproto_to_array(blob))[0]

"""
Explanation: Load mean values
End of explanation
"""


def prep_image(url, fname=None):
    if fname is None:
        ext = url.split('.')[-1]
        im = plt.imread(io.BytesIO(urllib.urlopen(url).read()), ext)
    else:
        ext = fname.split('.')[-1]
        im = plt.imread(fname, ext)
    h, w, _ = im.shape
    if h < w:
        im = skimage.transform.resize(im, (256, w*256/h), preserve_range=True)
    else:
        im = skimage.transform.resize(im, (h*256/w, 256), preserve_range=True)
    h, w, _ = im.shape
    im = im[h//2-112:h//2+112, w//2-112:w//2+112]
    rawim = np.copy(im).astype('uint8')
    im = np.swapaxes(np.swapaxes(im, 1, 2), 0, 1)
    im = im[::-1, :, :]
    im = im - mean_values
    return rawim, floatX(im[np.newaxis])

"""
Explanation: Image loader
End of explanation
"""


n = 5
m = 5
i = 0
for url in image_urls:
    print url
    try:
        rawim, im = prep_image(url)
    except:
        print 'Failed to download'
        continue

    prob_lasangne = np.array(lasagne.layers.get_output(net['prob'], im, deterministic=True).eval())[0]
    prob_caffe = net_caffe.forward_all(data=im)['prob'][0]

    
    print 'Lasagne:'
    res = sorted(zip(classes, prob_lasangne), key=lambda t: t[1], reverse=True)[:n]
    for c, p in res:
        print '  ', c, p
        
    print 'Caffe:'
    res = sorted(zip(classes, prob_caffe), key=lambda t: t[1], reverse=True)[:n]
    for c, p in res:
        print '  ', c, p
    
    plt.figure()
    plt.imshow(rawim.astype('uint8'))
    plt.axis('off')
    plt.show()
    
    i += 1
    if i == m:
        break
    
    print '\n\n'

model = {
    'values': lasagne.layers.get_all_param_values(net['prob']),
    'synset_words': classes,
    'mean_image': mean_values
}

pickle.dump(model, open('./resnet50.pkl', 'wb'), protocol=-1)

"""
Explanation: Lets take five images and compare prediction of Lasagne with Caffe
End of explanation
"""


