<gh_stars>10-100
"""
code borrow from:

https://github.com/mnicnc404/CartoonGan-tensorflow
https://github.com/xiaohu2015/DeepLearning_tutorials/blob/master/CNNs/shufflenet_v2.py
https://github.com/zengarden/light_head_rcnn
https://github.com/geonseoks/Light_head_R_CNN_xception
https://github.com/Stick-To/light-head-rcnn-tensorflow

"""

import tensorflow as tf #todo: remove this line
import tensorflow.keras.backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer, InputSpec

from tensorflow.keras.layers import Conv2D, DepthwiseConv2D, ReLU
from tensorflow.keras.layers import MaxPool2D, GlobalAveragePooling2D, Dense
from tensorflow.keras.layers import BatchNormalization, Activation

CEM_FILTER=245

@tf.function
def channle_shuffle(inputs, group):
    """Shuffle the channel
    Args:
        inputs: 4D Tensor
        group: int, number of groups
    Returns:
        Shuffled 4D Tensor
    """
    #in_shape = inputs.get_shape().as_list()
    h, w, in_channel  = K.int_shape(inputs)[1:]
    #h, w, in_channel = in_shape[1:]
    assert(in_channel % group == 0)
    l = K.reshape(inputs, [-1, h, w, in_channel // group, group])
    l = K.permute_dimensions(l, [0, 1, 2, 4, 3])
    l = K.reshape(l, [-1, h, w, in_channel])

    return l


class Conv2D_BN(Model):
    """Conv2D -> BN """
    def __init__(self, channel, kernel_size=1, stride=1):
        super(Conv2D_BN, self).__init__()

        self.conv = Conv2D(channel, kernel_size, strides=stride,
                            padding="SAME", use_bias=False)
        self.bn = BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)

    def build(self, input_shape):
        super(Conv2D_BN, self).build(input_shape)
    @tf.function
    def call(self, inputs, training=False):
        x = self.conv(inputs)
        x = self.bn(x, training=training)

        return x


class Conv2D_BN_ReLU(Model):
    """Conv2D -> BN -> ReLU"""
    def __init__(self, channel, kernel_size=1, stride=1):
        super(Conv2D_BN_ReLU, self).__init__(name="Conv2D_BN_ReLU")

        self.conv = Conv2D(channel, kernel_size, strides=stride,
                            padding="SAME", use_bias=False)
        self.bn_ = BatchNormalization(momentum=0.9, epsilon=1e-5)
        self.relu = Activation("relu")
        self.model = tf.keras.models.Sequential()
        self.model.add(self.conv)
        self.model.add(self.bn_ )
        self.model.add(self.relu)



    def build(self, input_shape):
        super(Conv2D_BN_ReLU, self).build(input_shape)

    @tf.function
    def call(self, inputs, training=False):
        #x = self.conv(inputs)
        #x = self.bn_(x, training=training)
        #x = self.relu(x)
        x=self.model(inputs, training=training)
        return x


class DepthwiseConv2D_BN(Model):
    """DepthwiseConv2D -> BN"""
    def __init__(self, kernel_size=3, stride=1):
        super(DepthwiseConv2D_BN, self).__init__()

        self.dconv = DepthwiseConv2D(kernel_size, strides=stride,
                                     depth_multiplier=1,
                                     padding="SAME", use_bias=False)
        self.bn = BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
    @tf.function
    def call(self, inputs, training=False):
        x = self.dconv(inputs)
        x = self.bn(x, training=training)
        return x

class DepthwiseConv2D_BN_POINT(Model):
    def __init__(self, kernel_size=3, stride=1, out_channel=256):
        super(DepthwiseConv2D_BN_POINT, self).__init__()
        self.dconv = DepthwiseConv2D(kernel_size, strides=stride,
                                     depth_multiplier=1,
                                     padding="SAME", use_bias=False)
        self.bn = BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
        self.conv = Conv2D(out_channel, 1, strides=1,
                            padding="SAME", use_bias=True)
    @tf.function
    def call(self, inputs, training=False):
        x = self.dconv(inputs)
        x = self.bn(x, training=training)
        x = self.conv(x)
        return x




class ShufflenetUnit1(Model):
    def __init__(self, out_channel):
        """The unit of shufflenetv2 for stride=1
        Args:
            out_channel: int, number of channels
        """
        super(ShufflenetUnit1, self).__init__()

        assert out_channel % 2 == 0
        self.out_channel = out_channel

        self.conv1_bn_relu = Conv2D_BN_ReLU(out_channel // 2, 1, 1)
        self.dconv_bn = DepthwiseConv2D_BN(5, 1)
        self.conv2_bn_relu = Conv2D_BN_ReLU(out_channel // 2, 1, 1)

    def build(self, input_shape):
        super(ShufflenetUnit1, self).build(input_shape)
    @tf.function
    def call(self, inputs, training=False):
        #print(K.int_shape(inputs))
        # split the channel
        shortcut, x = tf.split(inputs, 2, axis=3)

        x = self.conv1_bn_relu(x, training=training)
        x = self.dconv_bn(x, training=training)
        x = self.conv2_bn_relu(x, training=training)

        x = tf.concat([shortcut, x], axis=3)
        #print(K.int_shape(x))
        x = channle_shuffle(x, 2)
        return x

class ShufflenetUnit2(tf.keras.Model):
    """The unit of shufflenetv2 for stride=2"""
    def __init__(self, in_channel, out_channel):
        super(ShufflenetUnit2, self).__init__()

        assert out_channel % 2 == 0
        self.in_channel = in_channel
        self.out_channel = out_channel

        self.conv1_bn_relu = Conv2D_BN_ReLU(out_channel // 2, 1, 1)
        self.dconv_bn = DepthwiseConv2D_BN(5, 2)
        self.conv2_bn = Conv2D_BN(out_channel - in_channel, 1, 1)

        # for shortcut
        self.shortcut_dconv_bn = DepthwiseConv2D_BN(3, 2)
        self.shortcut_conv_bn = Conv2D_BN(in_channel, 1, 1)
    @tf.function
    def call(self, inputs, training=False):
        shortcut, x = inputs, inputs

        x = self.conv1_bn_relu(x, training=training)
        x = self.dconv_bn(x, training=training)
        x = self.conv2_bn(x, training=training)

        shortcut = self.shortcut_dconv_bn(shortcut, training=training)
        shortcut = self.shortcut_conv_bn(shortcut, training=training)

        x = tf.concat([shortcut, x], axis=3)
        x = ReLU()(x)
        x = channle_shuffle(x, 2)
        return x

class ShufflenetStage(Model):
    """The stage of shufflenet"""
    def __init__(self, in_channel, out_channel, num_blocks):
        super(ShufflenetStage, self).__init__()

        self.in_channel = in_channel
        self.out_channel = out_channel

        self.ops = []
        for i in range(num_blocks):
            if i == 0:
                op = ShufflenetUnit2(in_channel, out_channel)
            else:
                op = ShufflenetUnit1(out_channel)
            self.ops.append(op)
    @tf.function
    def call(self, inputs, training=False):
        x = inputs
        for op in self.ops:
            x = op(x, training=training)
        return x


class CEM(Model):
    """Context Enhancement Module"""
    def __init__(self):
        super(CEM, self).__init__()
        self.conv4 = Conv2D(CEM_FILTER, 1, strides=1,
                        padding="SAME", use_bias=True)
        self.conv5 = Conv2D(CEM_FILTER, 1, strides=1,
                        padding="SAME", use_bias=True)
        #self.b = K.reshape(inputs, [-1, h, w, in_channel // group, group])
    @tf.function
    def call(self, inputs, training=False):
        C4_lat = self.conv4(inputs[0])
        C5_lat = self.conv5(inputs[1])
        C5_lat = tf.keras.backend.resize_images(C5_lat, 2, 2, "channels_last", "bilinear")
        Cglb_lat = K.reshape(inputs[2], [-1, 1, 1, CEM_FILTER])
        return C4_lat + C5_lat + Cglb_lat


class ShuffleNetv2(Model):
    """Shufflenetv2"""
    def __init__(self, num_classes, first_channel=24, channels_per_stage=(132, 264, 528)):
        super(ShuffleNetv2, self).__init__(name="ShuffleNetv2")

        self.num_classes = num_classes

        self.conv1_bn_relu = Conv2D_BN_ReLU(first_channel, 3, 2)
        self.pool1 = MaxPool2D(3, strides=2, padding="SAME")
        self.stage2 = ShufflenetStage(first_channel, channels_per_stage[0], 4)
        self.stage3 = ShufflenetStage(channels_per_stage[0], channels_per_stage[1], 8)
        self.stage4 = ShufflenetStage(channels_per_stage[1], channels_per_stage[2], 4)
        #self.conv5_bn_relu = Conv2D_BN_ReLU(1024, 1, 1)
        self.gap = GlobalAveragePooling2D()
        self.linear = Dense(num_classes)

    def build(self, input_shape):
        super(ShuffleNetv2, self).build(input_shape)

    @tf.function
    def call(self, inputs, training=False):
        x = self.conv1_bn_relu(inputs, training=training)
        x = self.pool1(x)
        x = self.stage2(x, training=training)
        C4 = self.stage3(x, training=training)
        C5 = self.stage4(C4, training=training)
        #print("C5: ", K.int_shape(C5))
        #x = self.conv5_bn_relu(x, training=training)
        Cglb = self.gap(C5)
        print(K.int_shape(Cglb))
        x = self.linear(Cglb)
        #print(K.int_shape(x))

        return x, C4, C5, Cglb


class RPNROI(Model):
    """roi"""
    def __init__(self):
        super(RPNROI, self).__init()

    @tf.function
    def call(self, inputs, training=False):
        """inputs=[SAM, rpn_conf, rpn_pbbox] """
        pass #TODO

class RPN(Model):
    """region proposal network"""

    def __init__(self, filter=256):
        super(RPN, self).__init__()
        self.num_anchors = 5*5
        self.rpn = DepthwiseConv2D_BN_POINT(kernel_size=6, stride=1, out_channel=filter)
        self.rpn_cls_score = Conv2D(2*self.num_anchors, 1, strides=1,
                                padding="VALID", use_bias=True)
        self.rpn_cls_pred = Conv2D(4*self.num_anchors, 1, strides=1,
                                padding="VALID", use_bias=True)
    @tf.function
    def call(self, inputs, training=False):
        rpn = self.rpn(inputs, training=training)
        rpn_conf  = self.rpn_cls_score(rpn)
        #cls_pred  = tf.reshape(rpn_cls_score, [tf.shape(rpn_cls_score)[0], -1, 2]
        rpn_pbbox  = self.rpn_cls_pred(rpn)


        return rpn, rpn_conf, rpn_pbbox




class SAM(Model):
    """spatial attention module"""
    def __init__(self):
        super(SAM, self).__init__()
        self.point =  Conv2D(CEM_FILTER, 1, strides=1,
                                padding="VALID", use_bias=False)
        self.bn = BatchNormalization()

    @tf.function
    def call(self, inputs, training=False):
        """[RPN, CEM] """
        x = self.point(inputs[0])
        x = self.bn(x)
        x = tf.keras.activations.softmax(x, axis=-1)
        x = tf.math.multiply(x, inputs[1])
        return x



import numpy as np

if __name__ == "__main__":



    s = (10, 320, 320, 12)
    nx = np.random.rand(*s).astype(np.float32)

    custom_layers = [
        ShufflenetUnit1(out_channel = s[-1]),
        ShufflenetUnit2(in_channel = 24, out_channel = 116),
        ShufflenetStage(in_channel = 24, out_channel = 116, num_blocks = 5)
    ]

    for layer in custom_layers:
        tf.keras.backend.clear_session()
        out = layer(nx)
        layer.summary()
        print(f"Input  Shape: {nx.shape}")
        print(f"Output Shape: {out.shape}")
        print("\n" * 2)

    tf.keras.backend.clear_session()
    g = ShuffleNetv2(CEM_FILTER)





    shape = (10, 320, 320, 3)
    nx = np.random.rand(*shape).astype(np.float32)
    t = tf.keras.Input(shape=nx.shape[1:], batch_size=nx.shape[0])

    x,C4,C5,Cglb = g(nx, training=False)

    cem = CEM()
    sam = SAM()
    rpn = RPN()
    re = cem([C4, C5, x], training=False)

    rpn_result, rpn_cls_score, rpn_cls_pred = rpn(re, training =False)

    sam_result = sam([rpn_result, re], training=False)


    print('cem rsult: ', K.int_shape(re))


    g.summary()
    sam.summary()
    cem.summary()
    rpn.summary()


"""
code borrow from:

https://github.com/mnicnc404/CartoonGan-tensorflow
https://github.com/xiaohu2015/DeepLearning_tutorials/blob/master/CNNs/shufflenet_v2.py
https://github.com/zengarden/light_head_rcnn
https://github.com/geonseoks/Light_head_R_CNN_xception
https://github.com/Stick-To/light-head-rcnn-tensorflow

"""

import tensorflow as tf #todo: remove this line
import tensorflow.keras.backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Layer, InputSpec

from tensorflow.keras.layers import Conv2D, DepthwiseConv2D, ReLU
from tensorflow.keras.layers import MaxPool2D, GlobalAveragePooling2D, Dense
from tensorflow.keras.layers import BatchNormalization, Activation

CEM_FILTER=245

@tf.function
def channle_shuffle(inputs, group):
    """Shuffle the channel
    Args:
        inputs: 4D Tensor
        group: int, number of groups
    Returns:
        Shuffled 4D Tensor
    """
    #in_shape = inputs.get_shape().as_list()
    h, w, in_channel  = K.int_shape(inputs)[1:]
    #h, w, in_channel = in_shape[1:]
    assert(in_channel % group == 0)
    l = K.reshape(inputs, [-1, h, w, in_channel // group, group])
    l = K.permute_dimensions(l, [0, 1, 2, 4, 3])
    l = K.reshape(l, [-1, h, w, in_channel])

    return l


class Conv2D_BN(Model):
    """Conv2D -> BN """
    def __init__(self, channel, kernel_size=1, stride=1):
        super(Conv2D_BN, self).__init__()

        self.conv = Conv2D(channel, kernel_size, strides=stride,
                            padding="SAME", use_bias=False)
        self.bn = BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)

    def build(self, input_shape):
        super(Conv2D_BN, self).build(input_shape)
    @tf.function
    def call(self, inputs, training=False):
        x = self.conv(inputs)
        x = self.bn(x, training=training)

        return x


class Conv2D_BN_ReLU(Model):
    """Conv2D -> BN -> ReLU"""
    def __init__(self, channel, kernel_size=1, stride=1):
        super(Conv2D_BN_ReLU, self).__init__(name="Conv2D_BN_ReLU")

        self.conv = Conv2D(channel, kernel_size, strides=stride,
                            padding="SAME", use_bias=False)
        self.bn_ = BatchNormalization(momentum=0.9, epsilon=1e-5)
        self.relu = Activation("relu")
        self.model = tf.keras.models.Sequential()
        self.model.add(self.conv)
        self.model.add(self.bn_ )
        self.model.add(self.relu)



    def build(self, input_shape):
        super(Conv2D_BN_ReLU, self).build(input_shape)

    @tf.function
    def call(self, inputs, training=False):
        #x = self.conv(inputs)
        #x = self.bn_(x, training=training)
        #x = self.relu(x)
        x=self.model(inputs, training=training)
        return x


class DepthwiseConv2D_BN(Model):
    """DepthwiseConv2D -> BN"""
    def __init__(self, kernel_size=3, stride=1):
        super(DepthwiseConv2D_BN, self).__init__()

        self.dconv = DepthwiseConv2D(kernel_size, strides=stride,
                                     depth_multiplier=1,
                                     padding="SAME", use_bias=False)
        self.bn = BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
    @tf.function
    def call(self, inputs, training=False):
        x = self.dconv(inputs)
        x = self.bn(x, training=training)
        return x

class DepthwiseConv2D_BN_POINT(Model):
    def __init__(self, kernel_size=3, stride=1, out_channel=256):
        super(DepthwiseConv2D_BN_POINT, self).__init__()
        self.dconv = DepthwiseConv2D(kernel_size, strides=stride,
                                     depth_multiplier=1,
                                     padding="SAME", use_bias=False)
        self.bn = BatchNormalization(axis=-1, momentum=0.9, epsilon=1e-5)
        self.conv = Conv2D(out_channel, 1, strides=1,
                            padding="SAME", use_bias=True)
    @tf.function
    def call(self, inputs, training=False):
        x = self.dconv(inputs)
        x = self.bn(x, training=training)
        x = self.conv(x)
        return x




class ShufflenetUnit1(Model):
    def __init__(self, out_channel):
        """The unit of shufflenetv2 for stride=1
        Args:
            out_channel: int, number of channels
        """
        super(ShufflenetUnit1, self).__init__()

        assert out_channel % 2 == 0
        self.out_channel = out_channel

        self.conv1_bn_relu = Conv2D_BN_ReLU(out_channel // 2, 1, 1)
        self.dconv_bn = DepthwiseConv2D_BN(5, 1)
        self.conv2_bn_relu = Conv2D_BN_ReLU(out_channel // 2, 1, 1)

    def build(self, input_shape):
        super(ShufflenetUnit1, self).build(input_shape)
    @tf.function
    def call(self, inputs, training=False):
        #print(K.int_shape(inputs))
        # split the channel
        shortcut, x = tf.split(inputs, 2, axis=3)

        x = self.conv1_bn_relu(x, training=training)
        x = self.dconv_bn(x, training=training)
        x = self.conv2_bn_relu(x, training=training)

        x = tf.concat([shortcut, x], axis=3)
        #print(K.int_shape(x))
        x = channle_shuffle(x, 2)
        return x

class ShufflenetUnit2(tf.keras.Model):
    """The unit of shufflenetv2 for stride=2"""
    def __init__(self, in_channel, out_channel):
        super(ShufflenetUnit2, self).__init__()

        assert out_channel % 2 == 0
        self.in_channel = in_channel
        self.out_channel = out_channel

        self.conv1_bn_relu = Conv2D_BN_ReLU(out_channel // 2, 1, 1)
        self.dconv_bn = DepthwiseConv2D_BN(5, 2)
        self.conv2_bn = Conv2D_BN(out_channel - in_channel, 1, 1)

        # for shortcut
        self.shortcut_dconv_bn = DepthwiseConv2D_BN(3, 2)
        self.shortcut_conv_bn = Conv2D_BN(in_channel, 1, 1)
    @tf.function
    def call(self, inputs, training=False):
        shortcut, x = inputs, inputs

        x = self.conv1_bn_relu(x, training=training)
        x = self.dconv_bn(x, training=training)
        x = self.conv2_bn(x, training=training)

        shortcut = self.shortcut_dconv_bn(shortcut, training=training)
        shortcut = self.shortcut_conv_bn(shortcut, training=training)

        x = tf.concat([shortcut, x], axis=3)
        x = ReLU()(x)
        x = channle_shuffle(x, 2)
        return x

class ShufflenetStage(Model):
    """The stage of shufflenet"""
    def __init__(self, in_channel, out_channel, num_blocks):
        super(ShufflenetStage, self).__init__()

        self.in_channel = in_channel
        self.out_channel = out_channel

        self.ops = []
        for i in range(num_blocks):
            if i == 0:
                op = ShufflenetUnit2(in_channel, out_channel)
            else:
                op = ShufflenetUnit1(out_channel)
            self.ops.append(op)
    @tf.function
    def call(self, inputs, training=False):
        x = inputs
        for op in self.ops:
            x = op(x, training=training)
        return x


class CEM(Model):
    """Context Enhancement Module"""
    def __init__(self):
        super(CEM, self).__init__()
        self.conv4 = Conv2D(CEM_FILTER, 1, strides=1,
                        padding="SAME", use_bias=True)
        self.conv5 = Conv2D(CEM_FILTER, 1, strides=1,
                        padding="SAME", use_bias=True)
        #self.b = K.reshape(inputs, [-1, h, w, in_channel // group, group])
    @tf.function
    def call(self, inputs, training=False):
        C4_lat = self.conv4(inputs[0])
        C5_lat = self.conv5(inputs[1])
        C5_lat = tf.keras.backend.resize_images(C5_lat, 2, 2, "channels_last", "bilinear")
        Cglb_lat = K.reshape(inputs[2], [-1, 1, 1, CEM_FILTER])
        return C4_lat + C5_lat + Cglb_lat


class ShuffleNetv2(Model):
    """Shufflenetv2"""
    def __init__(self, num_classes, first_channel=24, channels_per_stage=(132, 264, 528)):
        super(ShuffleNetv2, self).__init__(name="ShuffleNetv2")

        self.num_classes = num_classes

        self.conv1_bn_relu = Conv2D_BN_ReLU(first_channel, 3, 2)
        self.pool1 = MaxPool2D(3, strides=2, padding="SAME")
        self.stage2 = ShufflenetStage(first_channel, channels_per_stage[0], 4)
        self.stage3 = ShufflenetStage(channels_per_stage[0], channels_per_stage[1], 8)
        self.stage4 = ShufflenetStage(channels_per_stage[1], channels_per_stage[2], 4)
        #self.conv5_bn_relu = Conv2D_BN_ReLU(1024, 1, 1)
        self.gap = GlobalAveragePooling2D()
        self.linear = Dense(num_classes)

    def build(self, input_shape):
        super(ShuffleNetv2, self).build(input_shape)

    @tf.function
    def call(self, inputs, training=False):
        x = self.conv1_bn_relu(inputs, training=training)
        x = self.pool1(x)
        x = self.stage2(x, training=training)
        C4 = self.stage3(x, training=training)
        C5 = self.stage4(C4, training=training)
        #print("C5: ", K.int_shape(C5))
        #x = self.conv5_bn_relu(x, training=training)
        Cglb = self.gap(C5)
        print(K.int_shape(Cglb))
        x = self.linear(Cglb)
        #print(K.int_shape(x))

        return x, C4, C5, Cglb


class RPNROI(Model):
    """roi"""
    def __init__(self):
        super(RPNROI, self).__init()

    @tf.function
    def call(self, inputs, training=False):
        """inputs=[SAM, rpn_conf, rpn_pbbox] """
        pass #TODO

class RPN(Model):
    """region proposal network"""

    def __init__(self, filter=256):
        super(RPN, self).__init__()
        self.num_anchors = 5*5
        self.rpn = DepthwiseConv2D_BN_POINT(kernel_size=6, stride=1, out_channel=filter)
        self.rpn_cls_score = Conv2D(2*self.num_anchors, 1, strides=1,
                                padding="VALID", use_bias=True)
        self.rpn_cls_pred = Conv2D(4*self.num_anchors, 1, strides=1,
                                padding="VALID", use_bias=True)
    @tf.function
    def call(self, inputs, training=False):
        rpn = self.rpn(inputs, training=training)
        rpn_conf  = self.rpn_cls_score(rpn)
        #cls_pred  = tf.reshape(rpn_cls_score, [tf.shape(rpn_cls_score)[0], -1, 2]
        rpn_pbbox  = self.rpn_cls_pred(rpn)


        return rpn, rpn_conf, rpn_pbbox




class SAM(Model):
    """spatial attention module"""
    def __init__(self):
        super(SAM, self).__init__()
        self.point =  Conv2D(CEM_FILTER, 1, strides=1,
                                padding="VALID", use_bias=False)
        self.bn = BatchNormalization()

    @tf.function
    def call(self, inputs, training=False):
        """[RPN, CEM] """
        x = self.point(inputs[0])
        x = self.bn(x)
        x = tf.keras.activations.softmax(x, axis=-1)
        x = tf.math.multiply(x, inputs[1])
        return x



import numpy as np

if __name__ == "__main__":



    s = (10, 320, 320, 12)
    nx = np.random.rand(*s).astype(np.float32)

    custom_layers = [
        ShufflenetUnit1(out_channel = s[-1]),
        ShufflenetUnit2(in_channel = 24, out_channel = 116),
        ShufflenetStage(in_channel = 24, out_channel = 116, num_blocks = 5)
    ]

    for layer in custom_layers:
        tf.keras.backend.clear_session()
        out = layer(nx)
        layer.summary()
        print(f"Input  Shape: {nx.shape}")
        print(f"Output Shape: {out.shape}")
        print("\n" * 2)

    tf.keras.backend.clear_session()
    g = ShuffleNetv2(CEM_FILTER)





    shape = (10, 320, 320, 3)
    nx = np.random.rand(*shape).astype(np.float32)
    t = tf.keras.Input(shape=nx.shape[1:], batch_size=nx.shape[0])

    x,C4,C5,Cglb = g(nx, training=False)

    cem = CEM()
    sam = SAM()
    rpn = RPN()
    re = cem([C4, C5, x], training=False)

    rpn_result, rpn_cls_score, rpn_cls_pred = rpn(re, training =False)

    sam_result = sam([rpn_result, re], training=False)


    print('cem rsult: ', K.int_shape(re))


    g.summary()
    sam.summary()
    cem.summary()
    rpn.summary()
code borrow from:  https://github.com/mnicnc404/CartoonGan-tensorflow https://github.com/xiaohu2015/DeepLearning_tutorials/blob/master/CNNs/shufflenet_v2.py https://github.com/zengarden/light_head_rcnn https://github.com/geonseoks/Light_head_R_CNN_xception https://github.com/Stick-To/light-head-rcnn-tensorflow #todo: remove this line Shuffle the channel     Args:         inputs: 4D Tensor         group: int, number of groups     Returns:         Shuffled 4D Tensor #in_shape = inputs.get_shape().as_list() #h, w, in_channel = in_shape[1:] Conv2D -> BN Conv2D -> BN -> ReLU #x = self.conv(inputs) #x = self.bn_(x, training=training) #x = self.relu(x) DepthwiseConv2D -> BN The unit of shufflenetv2 for stride=1         Args:             out_channel: int, number of channels #print(K.int_shape(inputs)) # split the channel #print(K.int_shape(x)) The unit of shufflenetv2 for stride=2 # for shortcut The stage of shufflenet Context Enhancement Module #self.b = K.reshape(inputs, [-1, h, w, in_channel // group, group]) Shufflenetv2 #self.conv5_bn_relu = Conv2D_BN_ReLU(1024, 1, 1) #print("C5: ", K.int_shape(C5)) #x = self.conv5_bn_relu(x, training=training) #print(K.int_shape(x)) roi inputs=[SAM, rpn_conf, rpn_pbbox] #TODO region proposal network #cls_pred  = tf.reshape(rpn_cls_score, [tf.shape(rpn_cls_score)[0], -1, 2] spatial attention module [RPN, CEM]