import unittest

import keras
import numpy as np
from keras import Model, Input
from keras.backend import random_normal
from keras.layers import Conv2D, MaxPooling2D, Add, UpSampling2D, Lambda, Activation, Reshape, TimeDistributed, Flatten, \
    Dense
from tensorflow import Tensor
from tensorflow.python import RandomNormal

from nets.classifier import RoiPoolingConv, get_resnet50_classifier
from nets.multiOutResNet import IFRCNN_Backbones as backbones, classifier_layers
from keras.layers import Concatenate

from nets.resnet import resnet50_classifier_layers, resnet50_classifier_layers_new
from utils.anchors import get_anchors


def IF_RPN(baselayers, num_anchors):
    """
    :param num_anchors: int, num box per anchor
    :param baselayers: List of feature maps [x0,x1]
    :return: rpn out
    :feature format: channel last
    """
    # get feature map
    assert len(baselayers) == 2, 'output of backbones must be list of 2 feature maps'
    x0 = baselayers[0]
    x1 = baselayers[1]
    print(x0.shape, x1.shape)
    """
    :param rpn_out : list [regression, classes]
    """
    rpn_out_0 = get_rpn_fusion(x0, num_anchors=num_anchors,count=0)
    rpn_out_1 = get_rpn_fusion(x1, num_anchors=num_anchors,count=1)
    """
    :returns fusion rpn output :List [regression, classes]
    """
    regression = Concatenate(axis=-2, name='classification')([rpn_out_0[0], rpn_out_1[0]])
    classes = Concatenate(axis=-2, name='regression')([rpn_out_0[1], rpn_out_1[1]])

    return [regression, classes]


def get_if_rcnn(inputs: Tensor):
    """
    :param inputs: Tensor from Input Layer
    :return:
    """
    # get back bone outputs
    if_backbones_out = backbones(inputs)

    return if_backbones_out

def get_rpn_fusion(base_layers, num_anchors,count=0):
    # ----------------------------------------------------#
    #   利用一个512通道的3x3卷积进行特征整合
    # ----------------------------------------------------#
    x2 = Conv2D(256,
                kernel_size=(3, 3),
                padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                ,name='rpn_conv1_'+str(count)+'_fusion2'
                )(base_layers)
    x3 = Conv2D(256, (5, 5), padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                ,name='rpn_conv1'+str(count)+'_fusion3'
                )(base_layers)
    x1 = Conv2D(256, (1, 1),
                padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                ,name='rpn_conv1'+str(count)+'_fusion1'
                )(base_layers)
    x = Concatenate(axis=-1)([x1, x2, x3])
    # ----------------------------------------------------#
    #   利用一个1x1卷积调整通道数，获得预测结果
    # ----------------------------------------------------#
    x_class = Conv2D(num_anchors, (1, 1), activation='sigmoid', kernel_initializer=RandomNormal(stddev=0.02)
                     ,name='rpn_out_class_'+str(count)
                     )(x)
    x_regr = Conv2D(num_anchors * 4, (1, 1), activation='linear', kernel_initializer=RandomNormal(stddev=0.02)
                    ,name='rpn_out_regress_'+str(count)
                    )(x)

    x_class = Reshape((-1, 1)
                      , name="classification_"+str(count)
                      )(x_class)
    x_regr = Reshape((-1, 4)
                     , name="regression_"+str(count)
                     )(x_regr)
    return [x_class, x_regr]

def get_if_classifier_(base_layers, input_rois, nb_classes=21, pooling_regions=None):
    # num_rois, 38, 38, 1024 -> num_rois, 14, 14, 2048
    # from frcnnproject.nets.RoiPoolingConv import PyramidROIAlign
    # out_roi_pool = PyramidROIAlign(pool_shape=pooling_regions)([input_rois] + base_layers)
    out_roi_pool0 = RoiPoolingConv(7)([base_layers[0], input_rois])
    out_roi_pool1 = RoiPoolingConv(7)([base_layers[1], input_rois])
    out_roi_pool = Concatenate(axis=2)([out_roi_pool0, out_roi_pool1])
    # print('OUT FPN roi pool shape{}, \n Out ROI pooling shape {}'.format(out_roi_pool.__class__,out_roi_pool1.__class__))
    # num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048
    out = classifier_layers(out_roi_pool)

    # num_rois, 1, 1, 1024 -> num_rois, 2048
    out = TimeDistributed(Flatten())(out)

    # num_rois, 1, 1, 1024 -> num_rois, nb_classes
    out_class = TimeDistributed(Dense(nb_classes, activation='softmax', kernel_initializer=RandomNormal(stddev=0.02)),
                                name='dense_class_{}'.format(nb_classes))(out)
    # num_rois, 1, 1, 1024 -> num_rois, 4 * (nb_classes-1)
    out_regr = TimeDistributed(
        Dense(4 * (nb_classes - 1), activation='linear', kernel_initializer=RandomNormal(stddev=0.02)),
        name='dense_regress_{}'.format(nb_classes))(out)
    return [out_class, out_regr]
def get_if_classifier_2(base_layers, input_rois, roi_size, num_classes=21,count=0):
    # batch_size, 38, 38, 1024 -> batch_size, num_rois, 14, 14, 1024
    # out_roi_pool0 = RoiPoolingConv(roi_size)([base_layers[0], input_rois])
    out_roi_pool = RoiPoolingConv(roi_size)([base_layers, input_rois])
    # out_roi_pool = Concatenate(axis=2)([out_roi_pool0, out_roi_pool1])

    # batch_size, num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048
    out = resnet50_classifier_layers_new(out_roi_pool,count=count)

    # batch_size, num_rois, 1, 1, 2048 -> batch_size, num_rois, 2048
    out = TimeDistributed(Flatten())(out)

    # batch_size, num_rois, 2048 -> batch_size, num_rois, num_classes
    out_class   = TimeDistributed(Dense(num_classes, activation='softmax',
                                        kernel_initializer=RandomNormal(stddev=0.02)),
                                  name='dense_class_{}'.format(num_classes)+str(count))(out)
    # batch_size, num_rois, 2048 -> batch_size, num_rois, 4 * (num_classes-1)
    out_regr    = TimeDistributed(Dense(4 * (num_classes - 1), activation='linear',
                                        kernel_initializer=RandomNormal(stddev=0.02)),
                                  name='dense_regress_{}'.format(num_classes)+str(count))(out)
    return [out_class, out_regr]


def get_model_if_2(num_classes, num_anchors=6):
    inputs = Input(shape=(None, None, 3))
    roi_input0 = Input(shape=(None, 4))
    roi_input1 = Input(shape=(None, 4))



    # ----------------------------------------------------#
    #   假设输入为600,600,3
    #   获得一个38,38,1024的共享特征层base_layers
    # ----------------------------------------------------#
    base_layers = backbones(inputs)
    # print(base_layers[0].shape,base_layers[1].shape)
    # ----------------------------------------------------#
    #   将共享特征层传入建议框网络
    #   该网络结果会对先验框进行调整获得建议框
    # ----------------------------------------------------#
    # rpn = IF_RPN(base_layers, num_anchors)
    rpn0 = get_rpn_fusion(base_layers[0],num_anchors,count=0)
    rpn1 = get_rpn_fusion(base_layers[1],num_anchors,count=1)
    classifier0 = get_if_classifier_2(base_layers[0],roi_input0,14,num_classes,count=0)
    classifier1 = get_if_classifier_2(base_layers[1], roi_input1, 14, num_classes,count=1)
    # ----------------------------------------------------#
    #   将共享特征层和建议框传入classifier网络
    #   该网络结果会对建议框进行调整获得预测框
    # ----------------------------------------------------#
    classifier = [Concatenate(axis=1,name='dense_class_{}'.format(num_classes))([classifier0[0],classifier1[0]]),
                  Concatenate(axis=1,name='dense_regress_{}'.format(num_classes))([classifier0[1], classifier1[1]])
                  ]

    model_rpn0 = Model(inputs, rpn0)
    model_rpn1 = Model(inputs, rpn1)

    model_all = Model([inputs, roi_input0,roi_input1], rpn0+rpn1 + classifier)
    return model_rpn0,model_rpn1, model_all
if __name__ == '__main__':
    model_rpn,model_rpn1, model_all = get_model_if_2(4)
    model_all.summary()


import unittest

import keras
import numpy as np
from keras import Model, Input
from keras.backend import random_normal
from keras.layers import Conv2D, MaxPooling2D, Add, UpSampling2D, Lambda, Activation, Reshape, TimeDistributed, Flatten, \
    Dense
from tensorflow import Tensor
from tensorflow.python import RandomNormal

from nets.classifier import RoiPoolingConv, get_resnet50_classifier
from nets.multiOutResNet import IFRCNN_Backbones as backbones, classifier_layers
from keras.layers import Concatenate

from nets.resnet import resnet50_classifier_layers, resnet50_classifier_layers_new
from utils.anchors import get_anchors


def IF_RPN(baselayers, num_anchors):
    """
    :param num_anchors: int, num box per anchor
    :param baselayers: List of feature maps [x0,x1]
    :return: rpn out
    :feature format: channel last
    """
    # get feature map
    assert len(baselayers) == 2, 'output of backbones must be list of 2 feature maps'
    x0 = baselayers[0]
    x1 = baselayers[1]
    print(x0.shape, x1.shape)
    """
    :param rpn_out : list [regression, classes]
    """
    rpn_out_0 = get_rpn_fusion(x0, num_anchors=num_anchors,count=0)
    rpn_out_1 = get_rpn_fusion(x1, num_anchors=num_anchors,count=1)
    """
    :returns fusion rpn output :List [regression, classes]
    """
    regression = Concatenate(axis=-2, name='classification')([rpn_out_0[0], rpn_out_1[0]])
    classes = Concatenate(axis=-2, name='regression')([rpn_out_0[1], rpn_out_1[1]])

    return [regression, classes]


def get_if_rcnn(inputs: Tensor):
    """
    :param inputs: Tensor from Input Layer
    :return:
    """
    # get back bone outputs
    if_backbones_out = backbones(inputs)

    return if_backbones_out

def get_rpn_fusion(base_layers, num_anchors,count=0):
    # ----------------------------------------------------#
    #   利用一个512通道的3x3卷积进行特征整合
    # ----------------------------------------------------#
    x2 = Conv2D(256,
                kernel_size=(3, 3),
                padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                ,name='rpn_conv1_'+str(count)+'_fusion2'
                )(base_layers)
    x3 = Conv2D(256, (5, 5), padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                ,name='rpn_conv1'+str(count)+'_fusion3'
                )(base_layers)
    x1 = Conv2D(256, (1, 1),
                padding='same',
                activation='relu',
                kernel_initializer=RandomNormal(stddev=0.02)
                ,name='rpn_conv1'+str(count)+'_fusion1'
                )(base_layers)
    x = Concatenate(axis=-1)([x1, x2, x3])
    # ----------------------------------------------------#
    #   利用一个1x1卷积调整通道数，获得预测结果
    # ----------------------------------------------------#
    x_class = Conv2D(num_anchors, (1, 1), activation='sigmoid', kernel_initializer=RandomNormal(stddev=0.02)
                     ,name='rpn_out_class_'+str(count)
                     )(x)
    x_regr = Conv2D(num_anchors * 4, (1, 1), activation='linear', kernel_initializer=RandomNormal(stddev=0.02)
                    ,name='rpn_out_regress_'+str(count)
                    )(x)

    x_class = Reshape((-1, 1)
                      , name="classification_"+str(count)
                      )(x_class)
    x_regr = Reshape((-1, 4)
                     , name="regression_"+str(count)
                     )(x_regr)
    return [x_class, x_regr]

def get_if_classifier_(base_layers, input_rois, nb_classes=21, pooling_regions=None):
    # num_rois, 38, 38, 1024 -> num_rois, 14, 14, 2048
    # from frcnnproject.nets.RoiPoolingConv import PyramidROIAlign
    # out_roi_pool = PyramidROIAlign(pool_shape=pooling_regions)([input_rois] + base_layers)
    out_roi_pool0 = RoiPoolingConv(7)([base_layers[0], input_rois])
    out_roi_pool1 = RoiPoolingConv(7)([base_layers[1], input_rois])
    out_roi_pool = Concatenate(axis=2)([out_roi_pool0, out_roi_pool1])
    # print('OUT FPN roi pool shape{}, \n Out ROI pooling shape {}'.format(out_roi_pool.__class__,out_roi_pool1.__class__))
    # num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048
    out = classifier_layers(out_roi_pool)

    # num_rois, 1, 1, 1024 -> num_rois, 2048
    out = TimeDistributed(Flatten())(out)

    # num_rois, 1, 1, 1024 -> num_rois, nb_classes
    out_class = TimeDistributed(Dense(nb_classes, activation='softmax', kernel_initializer=RandomNormal(stddev=0.02)),
                                name='dense_class_{}'.format(nb_classes))(out)
    # num_rois, 1, 1, 1024 -> num_rois, 4 * (nb_classes-1)
    out_regr = TimeDistributed(
        Dense(4 * (nb_classes - 1), activation='linear', kernel_initializer=RandomNormal(stddev=0.02)),
        name='dense_regress_{}'.format(nb_classes))(out)
    return [out_class, out_regr]
def get_if_classifier_2(base_layers, input_rois, roi_size, num_classes=21,count=0):
    # batch_size, 38, 38, 1024 -> batch_size, num_rois, 14, 14, 1024
    # out_roi_pool0 = RoiPoolingConv(roi_size)([base_layers[0], input_rois])
    out_roi_pool = RoiPoolingConv(roi_size)([base_layers, input_rois])
    # out_roi_pool = Concatenate(axis=2)([out_roi_pool0, out_roi_pool1])

    # batch_size, num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048
    out = resnet50_classifier_layers_new(out_roi_pool,count=count)

    # batch_size, num_rois, 1, 1, 2048 -> batch_size, num_rois, 2048
    out = TimeDistributed(Flatten())(out)

    # batch_size, num_rois, 2048 -> batch_size, num_rois, num_classes
    out_class   = TimeDistributed(Dense(num_classes, activation='softmax',
                                        kernel_initializer=RandomNormal(stddev=0.02)),
                                  name='dense_class_{}'.format(num_classes)+str(count))(out)
    # batch_size, num_rois, 2048 -> batch_size, num_rois, 4 * (num_classes-1)
    out_regr    = TimeDistributed(Dense(4 * (num_classes - 1), activation='linear',
                                        kernel_initializer=RandomNormal(stddev=0.02)),
                                  name='dense_regress_{}'.format(num_classes)+str(count))(out)
    return [out_class, out_regr]


def get_model_if_2(num_classes, num_anchors=6):
    inputs = Input(shape=(None, None, 3))
    roi_input0 = Input(shape=(None, 4))
    roi_input1 = Input(shape=(None, 4))



    # ----------------------------------------------------#
    #   假设输入为600,600,3
    #   获得一个38,38,1024的共享特征层base_layers
    # ----------------------------------------------------#
    base_layers = backbones(inputs)
    # print(base_layers[0].shape,base_layers[1].shape)
    # ----------------------------------------------------#
    #   将共享特征层传入建议框网络
    #   该网络结果会对先验框进行调整获得建议框
    # ----------------------------------------------------#
    # rpn = IF_RPN(base_layers, num_anchors)
    rpn0 = get_rpn_fusion(base_layers[0],num_anchors,count=0)
    rpn1 = get_rpn_fusion(base_layers[1],num_anchors,count=1)
    classifier0 = get_if_classifier_2(base_layers[0],roi_input0,14,num_classes,count=0)
    classifier1 = get_if_classifier_2(base_layers[1], roi_input1, 14, num_classes,count=1)
    # ----------------------------------------------------#
    #   将共享特征层和建议框传入classifier网络
    #   该网络结果会对建议框进行调整获得预测框
    # ----------------------------------------------------#
    classifier = [Concatenate(axis=1,name='dense_class_{}'.format(num_classes))([classifier0[0],classifier1[0]]),
                  Concatenate(axis=1,name='dense_regress_{}'.format(num_classes))([classifier0[1], classifier1[1]])
                  ]

    model_rpn0 = Model(inputs, rpn0)
    model_rpn1 = Model(inputs, rpn1)

    model_all = Model([inputs, roi_input0,roi_input1], rpn0+rpn1 + classifier)
    return model_rpn0,model_rpn1, model_all
if __name__ == '__main__':
    model_rpn,model_rpn1, model_all = get_model_if_2(4)
    model_all.summary()


:param num_anchors: int, num box per anchor     :param baselayers: List of feature maps [x0,x1]     :return: rpn out     :feature format: channel last # get feature map :param rpn_out : list [regression, classes] :returns fusion rpn output :List [regression, classes] :param inputs: Tensor from Input Layer     :return: # get back bone outputs # ----------------------------------------------------# #   利用一个512通道的3x3卷积进行特征整合 # ----------------------------------------------------# # ----------------------------------------------------# #   利用一个1x1卷积调整通道数，获得预测结果 # ----------------------------------------------------# # num_rois, 38, 38, 1024 -> num_rois, 14, 14, 2048 # from frcnnproject.nets.RoiPoolingConv import PyramidROIAlign # out_roi_pool = PyramidROIAlign(pool_shape=pooling_regions)([input_rois] + base_layers) # print('OUT FPN roi pool shape{}, \n Out ROI pooling shape {}'.format(out_roi_pool.__class__,out_roi_pool1.__class__)) # num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048 # num_rois, 1, 1, 1024 -> num_rois, 2048 # num_rois, 1, 1, 1024 -> num_rois, nb_classes # num_rois, 1, 1, 1024 -> num_rois, 4 * (nb_classes-1) # batch_size, 38, 38, 1024 -> batch_size, num_rois, 14, 14, 1024 # out_roi_pool0 = RoiPoolingConv(roi_size)([base_layers[0], input_rois]) # out_roi_pool = Concatenate(axis=2)([out_roi_pool0, out_roi_pool1]) # batch_size, num_rois, 14, 14, 1024 -> num_rois, 1, 1, 2048 # batch_size, num_rois, 1, 1, 2048 -> batch_size, num_rois, 2048 # batch_size, num_rois, 2048 -> batch_size, num_rois, num_classes # batch_size, num_rois, 2048 -> batch_size, num_rois, 4 * (num_classes-1) # ----------------------------------------------------# #   假设输入为600,600,3 #   获得一个38,38,1024的共享特征层base_layers # ----------------------------------------------------# # print(base_layers[0].shape,base_layers[1].shape) # ----------------------------------------------------# #   将共享特征层传入建议框网络 #   该网络结果会对先验框进行调整获得建议框 # ----------------------------------------------------# # rpn = IF_RPN(base_layers, num_anchors) # ----------------------------------------------------# #   将共享特征层和建议框传入classifier网络 #   该网络结果会对建议框进行调整获得预测框 # ----------------------------------------------------#