class AlexNet(object):
    '''
    input_x : 输入图片
    keep_prob : 全连接层drop out,训练时需要dropout，而在测试时需要将完整的网络因此不需要dropout
    num_class : 输出种类,不同的训练集网络差异在于Fc8层的输出神经元，因此我们将类别数设置为参数。
    skip_layer : 跳过的层，用于修改Alexnet
    weights_path: 权重加载的路径
    '''

    def __init__(self, input_x, keep_prob, num_classes, skip_layer, weights_path):
        self.input_x = input_x
        self.keep_prob = keep_prob
        self.num_classes = num_classes
        self.skip_layer = skip_layer
        if weights_path == 'Default':
            self.weights_path = 'bvlc_alexnet.npy'
        else:
            self.weights_path = weights_path
        # Create the AlexNet Network Define
        self.create()

    '''
    If padding == "SAME":
      output_spatial_shape[i] = ceil(input_spatial_shape[i] / strides[i])

    If padding == "VALID":
      output_spatial_shape[i] = ceil((input_spatial_shape[i] - (spatial_filter_shape[i]-1) * dilation_rate[i]) / strides[i])
    '''

    def conv(self, input, weightAndHeight, num_kernels, stride, name, padding='SAME', padding_num=0, groups=1):
        print('name is {} np.shape(input) {}'.format(name, np.shape(input)))

        input_channels = int(np.shape(input)[-1])
        '''
        pad (tensor ,paddings ,mode = 'CONSTANT' ,name = None ,constant_values = 0)
        此操作根据您指定的 paddings 来填充一个 tensor。paddings 是一个具有形状 [n, 2] 的整数张量，其中 n 是 tensor 的秩。
        对于每个输入维度 D，paddings [D, 0] 表示在该维度的 tensor 内容之前要添加多少个值，而 paddings[D, 1] 表示在该维度中的 tensor 内容之后要添加多少值。
        如果 mode 是 “REFLECT”，那么这两个paddings[D, 0] 和 paddings[D, 1] 不得大于 tensor.dim_size(D) - 1。
        如果 mode 是 “SYMMETRIC”，那么这两个 paddings[D, 0] 和 paddings[D, 1] 不得大于tensor.dim_size(D)。
        [
            [
                1, # 第0维，前面补1个(对应tensor的第一维)
                2  # 第0维，后面补2个
            ],
            [
                3, # 第1维，前面补3个(对应tensor的第二维)
                4  # 第2维，后面补4个
            ]
        ]
        input为四维，第一维是batch,第二维是Height,第三维是Weight，第四维是channel。
        '''
        if padding_num != 0:
            input = tf.pad(input, [[0, 0], [padding_num, padding_num], [padding_num, padding_num], [0, 0]])
        '''
        tf.nn.conv2d(input, filter, strides, padding, use_cudnn_on_gpu=None, name=None)
        input:指需要做卷积的输入图像，它要求是一个Tensor，具有[batch, in_height, in_width, in_channels]这样的shape,要求类型为float32和float64其中之一

        filter:相当于CNN中的卷积核，它要求是一个Tensor，具有[filter_height, filter_width, in_channels, out_channels]这样的shape

        strides：卷积时在图像每一维的步长，这是一个维度为4的向量
        '''
        convolve = lambda inputImg, filter: tf.nn.conv2d(inputImg, filter, strides=[1, stride, stride, 1],
                                                         padding=padding)

        with tf.variable_scope(name) as scope:
            weights = tf.get_variable('weights',
                                      shape=[weightAndHeight, weightAndHeight, input_channels / groups, num_kernels])
            biases = tf.get_variable('biases', shape=[num_kernels])

        if groups == 1:
            conv = convolve(input, weights)
        else:
            '''
            tf.split(value, num_or_size_splits,axis=0, num=None,name='split')
            axis:第几个维度
            '''
            input_groups = tf.split(value=input, num_or_size_splits=groups, axis=3)
            '''
            注意:
            input_groups和weights_groups的第三维不同
            '''
            weights_groups = tf.split(value=weights, num_or_size_splits=groups, axis=3)
            output_groups = [convolve(x, y) for x, y in zip(input_groups, weights_groups)]

            conv = tf.concat(axis=3, values=output_groups)

        withBias = tf.reshape(tf.nn.bias_add(conv, biases), conv.get_shape().as_list())
        relu = tf.nn.relu(withBias)
        return relu

    def maxPooling(self, input, filter_size, stride, name, padding='SAME'):
        print('name is {} np.shape(input) {}'.format(name, np.shape(input)))
        '''
        tf.nn.max_pool(value, ksize, strides, padding, name=None)
        第一个参数value：需要池化的输入，一般池化层接在卷积层后面，所以输入通常是feature map，依然是[batch, height, width, channels]这样的shape
        第二个参数ksize：池化窗口的大小，取一个四维向量，一般是[1, height, width, 1]，因为我们不想在batch和channels上做池化，所以这两个维度设为了1
        第三个参数strides：和卷积类似，窗口在每一个维度上滑动的步长，一般也是[1, stride,stride, 1]
        第四个参数padding：和卷积类似，可以取'VALID' 或者'SAME'
        返回一个Tensor，类型不变，shape仍然是[batch, height, width, channels]这种形式
        '''
        return tf.nn.max_pool(input, ksize=[1, filter_size, filter_size, 1], strides=[1, stride, stride, 1],
                              padding=padding, name=name)

    def lrn(self, input, radius, alpha, beta, name, bias=1.0):
        print('name is {} np.shape(input) {}'.format(name, np.shape(input)))
        '''
        tf.nn.local_response_normalization(input, depth_radius=None, bias=None, alpha=None, beta=None, name=None)
        局部响应归一化原理是仿造生物学上活跃的神经元对相邻神经元的抑制现象（侧抑制）:兴奋神经元抑制其邻居的能力
        因为ReLU神经元具有无限激活，我们需要LRN来规范化。 我们希望检测具有大响应的高频特征。 如果我们围绕兴奋神经元的局部邻域进行标准化，则与其邻居相比，它变得更加敏感。
        同时，它将抑制在任何给定的当地社区中均匀大的响应。 如果所有值都很大，那么对这些值进行标准化会减少所有这些值。
        '''
        return tf.nn.local_response_normalization(input=input, depth_radius=radius, bias=bias, alpha=alpha, beta=beta,
                                                  name=name)

    def fc(self, input, num_input, num_output, name, drop_out=0, relu=True):
        print('name is {} np.shape(input) {}'.format(name, np.shape(input)))
        with tf.variable_scope(name) as scope:
            weights = tf.get_variable('weights', shape=[num_input, num_output], trainable=True)
            biases = tf.get_variable('biases', shape=[num_output], trainable=True)

            # Linear
            act = tf.nn.xw_plus_b(input, weights, biases, name=scope.name)

            if relu == True:
                relu = tf.nn.relu(act)
                if drop_out == 0:
                    return relu
                else:
                    '''
                    With probability keep_prob, outputs the input element scaled up by 1 / keep_prob, otherwise outputs 0. The scaling is so that the expected sum is unchanged.' \
                    The scaling is so that the expected sum is unchanged.

                    The scaling factor is set to 1 / keep_prob, because the dropout is should be disabled at testing or evaluation.
                    During testing or evaluation, the activation of each unit is unscaled.
                    '''
                    return tf.nn.dropout(relu, 1.0 - drop_out)
            else:
                if drop_out == 0:
                    return act
                else:
                    return tf.nn.dropout(act, 1.0 - drop_out)

    def create(self):
        # layer1
        conv1 = self.conv(self.input_x, 11, 96, 4, name='conv1', padding='VALID')
        pool1 = self.maxPooling(conv1, filter_size=3, stride=2, name='pool1', padding='VALID')
        norm1 = self.lrn(pool1, 2, 2e-05, 0.75, name='norm1')

        # layer2
        conv2 = self.conv(norm1, 5, 256, 1, name='conv2', padding_num=2, padding='VALID', groups=2)
        pool2 = self.maxPooling(conv2, 3, 2, name='pool2', padding='VALID')
        norm2 = self.lrn(pool2, 2, 2e-05, 0.75, name='norm2')

        # layer3
        conv3 = self.conv(norm2, 3, 384, 1, name='conv3', padding_num=1, padding='VALID', groups=1)

        # layer4
        conv4 = self.conv(conv3, 3, 384, 1, name='conv4', padding_num=1, padding='VALID', groups=2)

        # layer 5
        conv5 = self.conv(conv4, 3, 256, 1, name='conv5', padding_num=1, padding='VALID', groups=2)
        pool5 = self.maxPooling(conv5, 3, 2, name='pool5', padding='VALID')

        # layer6
        flattened = tf.reshape(pool5, [-1, 6 * 6 * 256])
        fc6 = self.fc(input=flattened, num_input=6 * 6 * 256, num_output=4096, name='fc6', drop_out=1.0 - self.keep_prob,
                      relu=True)

        # layer 7
        fc7 = self.fc(input=fc6, num_input=4096, num_output=4096, name='fc7', drop_out=1.0 - self.keep_prob, relu=True)

        # layer 8
        self.fc8 = self.fc(input=fc7, num_input=4096, num_output=self.num_classes, name='fc8', drop_out=0, relu=False)

# 通过ImageNet预训练模型初始化我们的各层参数，各层以dict形式存在nparray中，因此我们首先加载模型，然后根据网络的名字分别加载权重及偏置的参数。
    # weights_dict['conv1']是一个键为weights和biases字典
    def load_weights(self, session):
        # 返回可遍历的(键, 值) 元组数组。
        weights_dict = np.load(self.weights_path, encoding='bytes').item()

        for op_name in weights_dict:
            if op_name not in self.skip_layer:
                with tf.variable_scope(op_name, reuse=True):
                    # Loop over list of weights/biases and assign them to their corresponding tf variable
                    for data in weights_dict[op_name]:
                        if len(data.shape) == 1:
                            var = tf.get_variable('biases', trainable=False)
                            session.run(var.assign(data))
                        else:
                            var = tf.get_variable('weights', trainable=False)
                            session.run(var.assign(data))

import cv2
import tensorflow as tf
import numpy as np
from Alexnet import AlexNet


def test_image(path_image, num_class, path_classes, weights_path='Default'):
    input = cv2.imread(path_image)
    input = cv2.resize(input, (227, 227))
    input = input.astype(np.float32)
    input = np.reshape(input, [1, 227, 227, 3])
    model = AlexNet(input, 0.9, 1000, skip_layer='', weights_path=weights_path)
    score = model.fc8
    max = tf.argmax(score, 1)
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        model.load_weights(sess)
        label_id = sess.run(max)[0]

        with open(path_classes) as f:
            lines = f.readlines()
            label = lines[label_id]
            print('image name is {} class_id is {} class_name is {}'.format(path_image, label_id, label))
        f.close()


test_image('/home/zyc/tensorflow-AlexNet/poodle.png', 1000, 'caffe_classes.py')