<reponame>JACKYLUO1991/Face-skin-hair-segmentaiton-and-skin-color-evaluation
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/3/27 17:42
# @Author  : JackyLUO
# @E-mail  : <EMAIL>
# @Site    : 
# @File    : lednet.py
# @Software: PyCharm

from keras import layers, models
import tensorflow as tf


class LEDNet:
    def __init__(self, groups, classes, input_shape):
        self.groups = groups
        self.classes = classes
        self.input_shape = input_shape

    def ss_bt(self, x, dilation, strides=(1, 1), padding='same'):
        x1, x2 = self.channel_split(x)
        filters = (int(x.shape[-1]) // self.groups)
        x1 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding)(x1)
        x1 = layers.Activation('relu')(x1)
        x1 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding)(x1)
        x1 = layers.BatchNormalization()(x1)
        x1 = layers.Activation('relu')(x1)
        x1 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding, dilation_rate=(dilation, 1))(
            x1)
        x1 = layers.Activation('relu')(x1)
        x1 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding, dilation_rate=(1, dilation))(
            x1)
        x1 = layers.BatchNormalization()(x1)
        x1 = layers.Activation('relu')(x1)

        x2 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding)(x2)
        x2 = layers.Activation('relu')(x2)
        x2 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding)(x2)
        x2 = layers.BatchNormalization()(x2)
        x2 = layers.Activation('relu')(x2)
        x2 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding, dilation_rate=(1, dilation))(
            x2)
        x2 = layers.Activation('relu')(x2)
        x2 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding, dilation_rate=(dilation, 1))(
            x2)
        x2 = layers.BatchNormalization()(x2)
        x2 = layers.Activation('relu')(x2)
        x_concat = layers.concatenate([x1, x2], axis=-1)
        x_add = layers.add([x, x_concat])
        output = self.channel_shuffle(x_add)
        return output

    def channel_shuffle(self, x):
        n, h, w, c = x.shape.as_list()
        x_reshaped = layers.Reshape([h, w, self.groups, int(c // self.groups)])(x)
        x_transposed = layers.Permute((1, 2, 4, 3))(x_reshaped)
        output = layers.Reshape([h, w, c])(x_transposed)
        return output

    def channel_split(self, x):
        def splitter(y):
            # keras Lambda saving bug!!!
            # x_left = layers.Lambda(lambda y: y[:, :, :, :int(int(y.shape[-1]) // self.groups)])(x)
            # x_right = layers.Lambda(lambda y: y[:, :, :, int(int(y.shape[-1]) // self.groups):])(x)
            # return x_left, x_right
            return tf.split(y, num_or_size_splits=self.groups, axis=-1)

        return layers.Lambda(lambda y: splitter(y))(x)

    def down_sample(self, x, filters):
        x_filters = int(x.shape[-1])
        x_conv = layers.Conv2D(filters - x_filters, kernel_size=3, strides=(2, 2), padding='same')(x)
        x_pool = layers.MaxPool2D()(x)
        x = layers.concatenate([x_conv, x_pool], axis=-1)
        x = layers.BatchNormalization()(x)
        x = layers.Activation('relu')(x)
        return x

    def apn_module(self, x):

        def right(x):
            x = layers.AveragePooling2D()(x)
            x = layers.Conv2D(self.classes, kernel_size=1, padding='same')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Activation('relu')(x)
            x = layers.UpSampling2D(interpolation='bilinear')(x)
            return x

        def conv(x, filters, kernel_size, stride):
            x = layers.Conv2D(filters, kernel_size=kernel_size, strides=(stride, stride), padding='same')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Activation('relu')(x)
            return x

        x_7 = conv(x, int(x.shape[-1]), 7, stride=2)
        x_5 = conv(x_7, int(x.shape[-1]), 5, stride=2)
        x_3 = conv(x_5, int(x.shape[-1]), 3, stride=2)

        x_3_1 = conv(x_3, self.classes, 3, stride=1)
        x_3_1_up = layers.UpSampling2D(interpolation='bilinear')(x_3_1)
        x_5_1 = conv(x_5, self.classes, 5, stride=1)
        x_3_5 = layers.add([x_5_1, x_3_1_up])
        x_3_5_up = layers.UpSampling2D(interpolation='bilinear')(x_3_5)
        x_7_1 = conv(x_7, self.classes, 3, stride=1)
        x_3_5_7 = layers.add([x_7_1, x_3_5_up])
        x_3_5_7_up = layers.UpSampling2D(interpolation='bilinear')(x_3_5_7)

        x_middle = conv(x, self.classes, 1, stride=1)
        x_middle = layers.multiply([x_3_5_7_up, x_middle])

        x_right = right(x)
        x_middle = layers.add([x_middle, x_right])
        return x_middle

    def encoder(self, x):
        x = self.down_sample(x, filters=32)
        for _ in range(3):
            x = self.ss_bt(x, dilation=1)

        x = self.down_sample(x, filters=64)
        for _ in range(2):
            x = self.ss_bt(x, dilation=1)

        x = self.down_sample(x, filters=128)

        dilation_rate = [1, 2, 5, 9, 2, 5, 9, 17]
        for dilation in dilation_rate:
            x = self.ss_bt(x, dilation=dilation)
        return x

    def decoder(self, x):
        x = self.apn_module(x)
        x = layers.UpSampling2D(size=8, interpolation='bilinear')(x)
        x = layers.Conv2D(self.classes, kernel_size=3, padding='same')(x)
        x = layers.BatchNormalization()(x)
        x = layers.Activation('softmax')(x)
        return x

    def model(self):
        inputs = layers.Input(shape=self.input_shape)
        encoder_out = self.encoder(inputs)
        output = self.decoder(encoder_out)
        return models.Model(inputs, output)


if __name__ == '__main__':
    from flops import get_flops

    model = LEDNet(2, 3, (256, 256, 3)).model()
    model.summary()

    get_flops(model)

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2020/3/27 17:42
# @Author  : JackyLUO
# @E-mail  : <EMAIL>
# @Site    : 
# @File    : lednet.py
# @Software: PyCharm

from keras import layers, models
import tensorflow as tf


class LEDNet:
    def __init__(self, groups, classes, input_shape):
        self.groups = groups
        self.classes = classes
        self.input_shape = input_shape

    def ss_bt(self, x, dilation, strides=(1, 1), padding='same'):
        x1, x2 = self.channel_split(x)
        filters = (int(x.shape[-1]) // self.groups)
        x1 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding)(x1)
        x1 = layers.Activation('relu')(x1)
        x1 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding)(x1)
        x1 = layers.BatchNormalization()(x1)
        x1 = layers.Activation('relu')(x1)
        x1 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding, dilation_rate=(dilation, 1))(
            x1)
        x1 = layers.Activation('relu')(x1)
        x1 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding, dilation_rate=(1, dilation))(
            x1)
        x1 = layers.BatchNormalization()(x1)
        x1 = layers.Activation('relu')(x1)

        x2 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding)(x2)
        x2 = layers.Activation('relu')(x2)
        x2 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding)(x2)
        x2 = layers.BatchNormalization()(x2)
        x2 = layers.Activation('relu')(x2)
        x2 = layers.Conv2D(filters, kernel_size=(1, 3), strides=strides, padding=padding, dilation_rate=(1, dilation))(
            x2)
        x2 = layers.Activation('relu')(x2)
        x2 = layers.Conv2D(filters, kernel_size=(3, 1), strides=strides, padding=padding, dilation_rate=(dilation, 1))(
            x2)
        x2 = layers.BatchNormalization()(x2)
        x2 = layers.Activation('relu')(x2)
        x_concat = layers.concatenate([x1, x2], axis=-1)
        x_add = layers.add([x, x_concat])
        output = self.channel_shuffle(x_add)
        return output

    def channel_shuffle(self, x):
        n, h, w, c = x.shape.as_list()
        x_reshaped = layers.Reshape([h, w, self.groups, int(c // self.groups)])(x)
        x_transposed = layers.Permute((1, 2, 4, 3))(x_reshaped)
        output = layers.Reshape([h, w, c])(x_transposed)
        return output

    def channel_split(self, x):
        def splitter(y):
            # keras Lambda saving bug!!!
            # x_left = layers.Lambda(lambda y: y[:, :, :, :int(int(y.shape[-1]) // self.groups)])(x)
            # x_right = layers.Lambda(lambda y: y[:, :, :, int(int(y.shape[-1]) // self.groups):])(x)
            # return x_left, x_right
            return tf.split(y, num_or_size_splits=self.groups, axis=-1)

        return layers.Lambda(lambda y: splitter(y))(x)

    def down_sample(self, x, filters):
        x_filters = int(x.shape[-1])
        x_conv = layers.Conv2D(filters - x_filters, kernel_size=3, strides=(2, 2), padding='same')(x)
        x_pool = layers.MaxPool2D()(x)
        x = layers.concatenate([x_conv, x_pool], axis=-1)
        x = layers.BatchNormalization()(x)
        x = layers.Activation('relu')(x)
        return x

    def apn_module(self, x):

        def right(x):
            x = layers.AveragePooling2D()(x)
            x = layers.Conv2D(self.classes, kernel_size=1, padding='same')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Activation('relu')(x)
            x = layers.UpSampling2D(interpolation='bilinear')(x)
            return x

        def conv(x, filters, kernel_size, stride):
            x = layers.Conv2D(filters, kernel_size=kernel_size, strides=(stride, stride), padding='same')(x)
            x = layers.BatchNormalization()(x)
            x = layers.Activation('relu')(x)
            return x

        x_7 = conv(x, int(x.shape[-1]), 7, stride=2)
        x_5 = conv(x_7, int(x.shape[-1]), 5, stride=2)
        x_3 = conv(x_5, int(x.shape[-1]), 3, stride=2)

        x_3_1 = conv(x_3, self.classes, 3, stride=1)
        x_3_1_up = layers.UpSampling2D(interpolation='bilinear')(x_3_1)
        x_5_1 = conv(x_5, self.classes, 5, stride=1)
        x_3_5 = layers.add([x_5_1, x_3_1_up])
        x_3_5_up = layers.UpSampling2D(interpolation='bilinear')(x_3_5)
        x_7_1 = conv(x_7, self.classes, 3, stride=1)
        x_3_5_7 = layers.add([x_7_1, x_3_5_up])
        x_3_5_7_up = layers.UpSampling2D(interpolation='bilinear')(x_3_5_7)

        x_middle = conv(x, self.classes, 1, stride=1)
        x_middle = layers.multiply([x_3_5_7_up, x_middle])

        x_right = right(x)
        x_middle = layers.add([x_middle, x_right])
        return x_middle

    def encoder(self, x):
        x = self.down_sample(x, filters=32)
        for _ in range(3):
            x = self.ss_bt(x, dilation=1)

        x = self.down_sample(x, filters=64)
        for _ in range(2):
            x = self.ss_bt(x, dilation=1)

        x = self.down_sample(x, filters=128)

        dilation_rate = [1, 2, 5, 9, 2, 5, 9, 17]
        for dilation in dilation_rate:
            x = self.ss_bt(x, dilation=dilation)
        return x

    def decoder(self, x):
        x = self.apn_module(x)
        x = layers.UpSampling2D(size=8, interpolation='bilinear')(x)
        x = layers.Conv2D(self.classes, kernel_size=3, padding='same')(x)
        x = layers.BatchNormalization()(x)
        x = layers.Activation('softmax')(x)
        return x

    def model(self):
        inputs = layers.Input(shape=self.input_shape)
        encoder_out = self.encoder(inputs)
        output = self.decoder(encoder_out)
        return models.Model(inputs, output)


if __name__ == '__main__':
    from flops import get_flops

    model = LEDNet(2, 3, (256, 256, 3)).model()
    model.summary()

    get_flops(model)
#!/usr/bin/env python # -*- coding: utf-8 -*- # @Time    : 2020/3/27 17:42 # @Author  : JackyLUO # @E-mail  : <EMAIL> # @Site    : # @File    : lednet.py # @Software: PyCharm # keras Lambda saving bug!!! # x_left = layers.Lambda(lambda y: y[:, :, :, :int(int(y.shape[-1]) // self.groups)])(x) # x_right = layers.Lambda(lambda y: y[:, :, :, int(int(y.shape[-1]) // self.groups):])(x) # return x_left, x_right