import torch
import torch.nn as nn
import torchvision
from core.config import cfg
from . import resnet, resnext,senet
import modeling.ResNet as ResNet
import modeling.ResNet_train as ResNet_train
import modeling.ResNet_test as ResNet_test
import modeling.deform_conv as deform_conv
from lib.nn import SynchronizedBatchNorm2d
import math
#from modeling.backbone import build_backbone
#from modeling.aspp import build_aspp
#from modeling.cspn import Affinity_Propagate
#from modeling.spn_bat import SPN
import utils.resnet_weights_helper as resnet_utils
import pickle
def get_func(func_name):
    """Helper to return a function object by name. func_name must identify a
    function in this module or the path to a function relative to the base
    'modeling' module.
    """
    if func_name == '':
        return None
    try:
        parts = func_name.split('.')
        # Refers to a function in this module
        if len(parts) == 1:
            return globals()[parts[0]]
        # Otherwise, assume we're referencing a module under modeling
        module_name = 'modeling.' + '.'.join(parts[:-1])
        module = importlib.import_module(module_name)
        return getattr(module, parts[-1])
    except Exception:
        logger.error('Failed to find function: %s', func_name)
        raise

class SegmentationModuleBase(nn.Module):
    def __init__(self):
        super(SegmentationModuleBase, self).__init__()

    def pixel_acc(self, pred, label):
        _, preds = torch.max(pred, dim=1)
        valid = (label >= 0).long()
        acc_sum = torch.sum(valid * (preds == label).long())
        pixel_sum = torch.sum(valid)
        acc = acc_sum.float() / (pixel_sum.float() + 1e-10)
        return acc


class SegmentationModule(SegmentationModuleBase):
    def __init__(self, net_enc, net_dec, crit, deep_sup_scale=None):
        super(SegmentationModule, self).__init__()
        self.encoder = net_enc
        self.decoder = net_dec
        self.crit = crit
        self.deep_sup_scale = deep_sup_scale

    def forward(self, feed_dict, *, segSize=None):
        if segSize is None: # training
            if self.deep_sup_scale is not None: # use deep supervision technique
                (pred, pred_deepsup) = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True))
            else:
                pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True))

            loss = self.crit(pred, feed_dict['seg_label'])
            if self.deep_sup_scale is not None:
                loss_deepsup = self.crit(pred_deepsup, feed_dict['seg_label'])
                loss = loss + loss_deepsup * self.deep_sup_scale

            acc = self.pixel_acc(pred, feed_dict['seg_label'])
            return loss, acc
        else: # inference
            pred = self.decoder(self.encoder(feed_dict['img_data'], return_feature_maps=True), segSize=cfg.SEM.INPUT_SIZE)
            return pred


def conv3x3(in_planes, out_planes, stride=1, has_bias=False):
    "3x3 convolution with padding"
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
            padding=1, bias=has_bias, groups=cfg.RESNETS.NUM_GROUPS)
            #padding=1, bias=has_bias, groups=cfg.RESNETS.NUM_GROUPS if cfg.SEM.DECODER_TYPE != 'deeplab' else 1)

def conv3x3_bn_relu(in_planes, out_planes, stride=1):
    return nn.Sequential(
            conv3x3(in_planes, out_planes, stride),
            SynchronizedBatchNorm2d(out_planes),
            nn.ReLU(inplace=True),
            )
class GCN(nn.Module):
    def __init__(self, in_planes, out_planes, kernel=23, stride=1, has_bias=False):
        super(GCN, self).__init__()
        #self.psp_conv = conv3x3_bn_relu(in_planes, out_planes, 1)
        #in_planes = out_planes
        #out_planes = cfg.MODEL.NUM_CLASSES
        self.gcn23x23_0 = nn.Sequential(
                nn.Conv2d(in_planes, out_planes, kernel_size=[1,kernel], stride=stride,
                padding=[0, kernel//2], dilation=(1,1), bias=has_bias, groups=cfg.RESNETS.NUM_GROUPS),
                SynchronizedBatchNorm2d(out_planes),
                nn.ReLU(inplace=True))

        self.gcn23x23_1 = nn.Sequential(
                nn.Conv2d(in_planes, out_planes, kernel_size=[kernel,1], stride=stride,
                padding=[kernel//2, 0], bias=has_bias, groups=cfg.RESNETS.NUM_GROUPS),
                SynchronizedBatchNorm2d(out_planes),
                nn.ReLU(inplace=True))
                
        self.gcn23x23_2 = nn.Sequential(
                nn.Conv2d(out_planes, out_planes, kernel_size=[kernel,1], stride=stride,
                padding=[kernel//2, 0], bias=has_bias, groups=cfg.RESNETS.NUM_GROUPS),
                SynchronizedBatchNorm2d(out_planes),
                nn.ReLU(inplace=True))

        self.gcn23x23_3 = nn.Sequential(
                nn.Conv2d(out_planes, out_planes, kernel_size=[1,kernel], stride=stride,
                padding=[0, kernel//2], bias=has_bias, groups=cfg.RESNETS.NUM_GROUPS),
                SynchronizedBatchNorm2d(out_planes),
                nn.ReLU(inplace=True))

        #self.br_side = nn.Sequential(
        #        nn.Conv2d(out_planes, cfg.MODEL.NUM_CLASSES, 3, padding=1, stride=1), 
        #        SynchronizedBatchNorm2d(cfg.MODEL.NUM_CLASSES),
        #        nn.ReLU(inplace=True),
        #        nn.Conv2d(cfg.MODEL.NUM_CLASSES, cfg.MODEL.NUM_CLASSES, 3, padding=1, stride=1))
        #self.conv_last = nn.Conv2d(out_planes, cfg.MODEL.NUM_CLASSES, 1, padding=0, stride=1)

        self.apply(self.weights_init)

    def weights_init(self, m):
        # custom weights initialization
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            nn.init.kaiming_normal_(m.weight.data)
        elif classname.find('BatchNorm') != -1:
            m.weight.data.fill_(1.)
            m.bias.data.fill_(1e-4)

    def forward(self, x):
        #x = self.psp_conv(x)
        out0 = self.gcn23x23_0(x)
        out2 = self.gcn23x23_2(out0)

        out1 = self.gcn23x23_1(x)
        out3 = self.gcn23x23_3(out1)
        
        out = out2 + out3
        #out_br = self.br_side(out)
        #out = self.conv_last(out)
        return out

class ModelBuilder():
    # custom weights initialization
    def weights_init(self, m):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            nn.init.kaiming_normal_(m.weight.data)
        elif classname.find('BatchNorm') != -1:
            m.weight.data.fill_(1.)
            m.bias.data.fill_(1e-4)
        #elif classname.find('Linear') != -1:
        #    m.weight.data.normal_(0.0, 0.0001)
        
    def build_encoder(self, arch='resnet50_dilated8', fc_dim=512, weights=''):
        if cfg.SEM.FREEZE_BN:
            print("Using AffineChannel2d as SynchronizedBatchNorm2d")
            from lib.nn import AffineChannel2d
            global SynchronizedBatchNorm2d
            SynchronizedBatchNorm2d = AffineChannel2d
        pretrained = True if len(weights) == 0 else False
        if arch == 'resnet18':
            orig_resnet = resnet.__dict__['resnet18'](pretrained=pretrained)
            net_encoder = Resnet(orig_resnet)
        elif arch == 'resnet18_dilated8':
            orig_resnet = resnet.__dict__['resnet18'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=8)
        elif arch == 'resnet18_dilated16':
            orig_resnet = resnet.__dict__['resnet18'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=16)
        elif arch == 'resnet34':
            raise NotImplementedError
            orig_resnet = resnet.__dict__['resnet34'](pretrained=pretrained)
            net_encoder = Resnet(orig_resnet)
        elif arch == 'resnet34_dilated8':
            raise NotImplementedError
            orig_resnet = resnet.__dict__['resnet34'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=8)
        elif arch == 'resnet34_dilated16':
            raise NotImplementedError
            orig_resnet = resnet.__dict__['resnet34'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=16)
        elif arch == 'resnet50':
            orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained)
            net_encoder = Resnet(orig_resnet)

        elif arch == 'ResnetX_dilated8' or arch == 'ResnetX_dilated16':
            #orig_resnet = ResNet.__dict__['resnet101'](pretrained=False)
            net_encoder = eval(cfg.MODEL.CONV_BODY)()
            print ('loading weights for ResNet')
            resnet_utils.load_pretrained_imagenet_weights(net_encoder)
            print ('loading weights is done')
            if cfg.TRAIN.FREEZE_CONV_BODY:
                print ('freeze train conv_body')
                for p in net_encoder.parameters():
                    p.requires_grad = False

            #print ('loading pretrained model for ResNet')
            #pretrained=torch.load(cfg.RESNETS.IMAGENET_PRETRAINED_WEIGHTS, map_location=lambda storage, loc: storage)
            #with open(cfg.RESNETS.IMAGENET_PRETRAINED_WEIGHTS, 'rb') as f:
            #    pretrained = pickle.load(f, encoding='latin1')
            #print (pretrained['blobs'].keys())
            #orig_resnet.load_state_dict(pretrained['model'],strict=True)
            #print (pretrained.keys())
            #net_encoder.load_state_dict(pretrained,strict=True)
            #print ('loading pretrained is done')
            #net_encoder = ResnetDilated(orig_resnet,dilate_scale=8)
            #net_encoder = ResnetDilated(orig_resnet,dilate_scale=8)
            ##pretrained=pretrained['model']
           
            #print ('loading pretrained is done')
            #from  lib.nn import AffineChannel2d 
            #global SynchronizedBatchNorm2d 
            #SynchronizedBatchNorm2d = AffineChannel2d

        elif arch == 'resnet50_dilated8':
            orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=8)

        elif arch == 'resnet50_dilated16':
            orig_resnet = resnet.__dict__['resnet50'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=16)
        elif arch == 'resnet101':
            orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
            net_encoder = Resnet(orig_resnet)
        elif arch == 'resnet101_dilated8':
            orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=8)
        elif arch == 'resnet101_dilated16':
            orig_resnet = resnet.__dict__['resnet101'](pretrained=pretrained)
            net_encoder = ResnetDilated(orig_resnet,
                                        dilate_scale=16)
        
        elif arch == 'resnet152':
            orig_resnet = resnet.__dict__['resnet152'](pretrained=pretrained)
            net_encoder = Resnet152(orig_resnet)
        elif arch == 'resnet152_dilated8':
            orig_resnet = resnet.__dict__['resnet152'](pretrained=pretrained)
            net_encoder = Resnet152Dilated(orig_resnet,
                                        dilate_scale=8)
        elif arch == 'resnet152_dilated16':
            orig_resnet = resnet.__dict__['resnet152'](pretrained=pretrained)
            net_encoder = Resnet152Dilated(orig_resnet,
                                        dilate_scale=16)
        
        elif arch == 'resnext101':
            orig_resnext = resnext.__dict__['resnext101'](pretrained=pretrained)
            net_encoder = Resnet(orig_resnext) # we can still use class Resnet
        
        elif arch == 'se_resnet50_dilate8':
            net_encoder = senet.__dict__['se_resnet50_dilate'](dilate=8,pretrained=pretrained)
        elif arch == 'se_resnet101_dilate8':
            net_encoder = senet.__dict__['se_resnet101_dilate'](dilate=8,pretrained=pretrained)
        elif arch == 'se_resnet152_dilate8':
            net_encoder = senet.__dict__['se_resnet152_dilate'](dilate=8,pretrained=pretrained)
        elif arch == 'se_resnext50_dilate8_32x4d':
            net_encoder = senet.__dict__['se_resnext50_dilate_32x4d'](dilate=8,pretrained=pretrained)
        elif arch == 'se_resnext101_dilate8_32x4d':
            net_encoder = senet.__dict__['se_resnext101_dilate_32x4d'](dilate=8,pretrained=pretrained)
        elif arch == 'senet154_dilate8':
            net_encoder = senet.__dict__['senet154_dilate'](dilate=8,pretrained=pretrained)

        elif arch == 'se_resnet50_dilate16':
            net_encoder = senet.__dict__['se_resnet50_dilate'](dilate=16,pretrained=pretrained)
        elif arch == 'se_resnet101_dilate16':
            net_encoder = senet.__dict__['se_resnet101_dilate'](dilate=16,pretrained=pretrained)
        elif arch == 'se_resnet152_dilate16':
            net_encoder = senet.__dict__['se_resnet152_dilate'](dilate=16,pretrained=pretrained)
        elif arch == 'se_resnext50_dilate16_32x4d':
            net_encoder = senet.__dict__['se_resnext50_dilate_32x4d'](dilate=16,pretrained=pretrained)
        elif arch == 'se_resnext101_dilate16_32x4d':
            net_encoder = senet.__dict__['se_resnext101_dilate_32x4d'](dilate=16,pretrained=pretrained)
        elif arch == 'senet154_dilate16':
            net_encoder = senet.__dict__['senet154_dilate'](dilate=16,pretrained=False)
        elif arch == 'xception':
            net_encoder=build_backbone('xception', 16, BatchNorm)
        else:
            raise Exception('Architecture undefined!')

        # net_encoder.apply(self.weights_init)
        if len(weights) > 0:
            print('Loading weights for net_encoder')
            net_encoder.load_state_dict(
                torch.load(weights, map_location=lambda storage, loc: storage), strict=False)
        return net_encoder

    def build_decoder(self, arch='ppm_bilinear_deepsup',
                      fc_dim=512, num_class=150,
                      weights='', use_softmax=False):
        if arch == 'c1_bilinear_deepsup':
            net_decoder = C1BilinearDeepSup(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax)
        elif arch == 'c1_bilinear':
            net_decoder = C1Bilinear(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax)
        elif arch == 'ppm_bilinear':
            net_decoder = PPMBilinear(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax)
        elif arch == 'ppm_bilinear_deepsup':
            net_decoder = PPMBilinearDeepsup(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax)
        elif arch == 'upernet_lite':
            net_decoder = UPerNet(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax,
                fpn_dim=256)
        elif arch == 'upernet':
            net_decoder = UPerNet(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax,
                fpn_dim=512)
        elif arch == 'cspn_upernet':
            net_decoder = CspnUPerNet(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax,
                fpn_dim=512)
        elif arch == 'spn_upernet':
            net_decoder = SpnUPerNet(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax,
                fpn_dim=512)
        elif arch == 'upernet_tmp':
            net_decoder = UPerNetTmp(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax,
                fpn_dim=512)
        elif arch == 'deeplab':
            net_decoder = DeepLabNet(
                num_class=num_class,
                fc_dim=fc_dim,
                use_softmax=use_softmax,
                fpn_dim=256,
                fpn_inplanes=(256,2048),
                pool_scales=(1,6,12,18) if '16' in cfg.SEM.ARCH_ENCODER else (1, 12, 24, 36))
        else:
            raise Exception('Architecture undefined!')

        net_decoder.apply(self.weights_init)
        if len(weights) > 0:
            print('Loading weights for net_decoder')
            net_decoder.load_state_dict(
                torch.load(weights, map_location=lambda storage, loc: storage), strict=False)
        return net_decoder


class Resnet(nn.Module):
    def __init__(self, orig_resnet):
        super(Resnet, self).__init__()

        # take pretrained resnet, except AvgPool and FC
        self.conv1 = orig_resnet.conv1
        self.bn1 = orig_resnet.bn1
        self.relu1 = orig_resnet.relu1
        self.conv2 = orig_resnet.conv2
        self.bn2 = orig_resnet.bn2
        self.relu2 = orig_resnet.relu2
        self.conv3 = orig_resnet.conv3
        self.bn3 = orig_resnet.bn3
        self.relu3 = orig_resnet.relu3
        self.maxpool = orig_resnet.maxpool
        self.layer1 = orig_resnet.layer1
        self.layer2 = orig_resnet.layer2
        self.layer3 = orig_resnet.layer3
        self.layer4 = orig_resnet.layer4

    def forward(self, x, return_feature_maps=False):
        conv_out = []

        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        x = self.relu3(self.bn3(self.conv3(x)))
        x = self.maxpool(x)

        x = self.layer1(x); conv_out.append(x);
        x = self.layer2(x); conv_out.append(x);
        x = self.layer3(x); conv_out.append(x);
        x = self.layer4(x); conv_out.append(x);

        if return_feature_maps:
            return conv_out
        return [x]


class ResnetDilated(nn.Module):
    def __init__(self, orig_resnet, dilate_scale=8):
        super(ResnetDilated, self).__init__()
        from functools import partial

        if dilate_scale == 8:
            orig_resnet.layer3.apply(
                partial(self._nostride_dilate, dilate=2))
            orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=4))
        elif dilate_scale == 16:
            orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=2))

        # take pretrained resnet, except AvgPool and FC
        self.conv1 = orig_resnet.conv1
        self.bn1 = orig_resnet.bn1
        self.relu1 = orig_resnet.relu1
        self.conv2 = orig_resnet.conv2
        self.bn2 = orig_resnet.bn2
        self.relu2 = orig_resnet.relu2
        self.conv3 = orig_resnet.conv3
        self.bn3 = orig_resnet.bn3
        self.relu3 = orig_resnet.relu3
        self.maxpool = orig_resnet.maxpool
        self.layer1 = orig_resnet.layer1
        self.layer2 = orig_resnet.layer2
        self.layer3 = orig_resnet.layer3
        self.layer4 = orig_resnet.layer4

    def _nostride_dilate(self, m, dilate):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            # the convolution with stride
            if m.stride == (2, 2):
                m.stride = (1, 1)
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate//2, dilate//2)
                    m.padding = (dilate//2, dilate//2)
            # other convoluions
            else:
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate, dilate)
                    m.padding = (dilate, dilate)

    def forward(self, x, return_feature_maps=False):
        conv_out = []

        x = self.relu1(self.bn1(self.conv1(x)))
        x = self.relu2(self.bn2(self.conv2(x)))
        x = self.relu3(self.bn3(self.conv3(x)))
        x = self.maxpool(x)

        x = self.layer1(x); conv_out.append(x);
        x = self.layer2(x); conv_out.append(x);
        x = self.layer3(x); conv_out.append(x);
        x = self.layer4(x); conv_out.append(x);

        if return_feature_maps:
            return conv_out
        return [x]

class Resnet152Dilated(nn.Module):
    def __init__(self, orig_resnet, dilate_scale=8):
        super(Resnet152Dilated, self).__init__()
        from functools import partial

        if dilate_scale == 8:
            orig_resnet.layer3.apply(
                partial(self._nostride_dilate, dilate=2))
            orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=4))
        elif dilate_scale == 16:
            orig_resnet.layer4.apply(
                partial(self._nostride_dilate, dilate=2))

        # take pretrained resnet, except AvgPool and FC
        self.conv1 = orig_resnet.conv1
        self.bn1 = orig_resnet.bn1
        self.relu = orig_resnet.relu
        self.maxpool = orig_resnet.maxpool
        self.layer1 = orig_resnet.layer1
        self.layer2 = orig_resnet.layer2
        self.layer3 = orig_resnet.layer3
        self.layer4 = orig_resnet.layer4

    def _nostride_dilate(self, m, dilate):
        classname = m.__class__.__name__
        if classname.find('Conv') != -1:
            # the convolution with stride
            if m.stride == (2, 2):
                m.stride = (1, 1)
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate//2, dilate//2)
                    m.padding = (dilate//2, dilate//2)
            # other convoluions
            else:
                if m.kernel_size == (3, 3):
                    m.dilation = (dilate, dilate)
                    m.padding = (dilate, dilate)

    def forward(self, x, return_feature_maps=False):
        conv_out = []

        x = self.relu(self.bn1(self.conv1(x)))
        x = self.maxpool(x)

        x = self.layer1(x); conv_out.append(x);
        x = self.layer2(x); conv_out.append(x);
        x = self.layer3(x); conv_out.append(x);
        x = self.layer4(x); conv_out.append(x);

        if return_feature_maps:
            return conv_out
        return [x]

# last conv, bilinear upsample
class C1BilinearDeepSup(nn.Module):
    def __init__(self, num_class=150, fc_dim=2048, use_softmax=False):
        super(C1BilinearDeepSup, self).__init__()
        self.use_softmax = use_softmax

        self.cbr = conv3x3_bn_relu(fc_dim, fc_dim // 4, 1)
        self.cbr_deepsup = conv3x3_bn_relu(fc_dim // 2, fc_dim // 4, 1)

        # last conv
        self.conv_last = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)
        self.conv_last_deepsup = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)

    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        x = self.cbr(conv5)
        x = self.conv_last(x)

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        # deep sup
        conv4 = conv_out[-2]
        _ = self.cbr_deepsup(conv4)
        _ = self.conv_last_deepsup(_)

        x = nn.functional.log_softmax(x, dim=1)
        _ = nn.functional.log_softmax(_, dim=1)

        return (x, _)


# last conv, bilinear upsample
class C1Bilinear(nn.Module):
    def __init__(self, num_class=150, fc_dim=2048, use_softmax=False):
        super(C1Bilinear, self).__init__()
        self.use_softmax = use_softmax

        self.cbr = conv3x3_bn_relu(fc_dim, fc_dim // 4, 1)

        # last conv
        self.conv_last = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)

    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]
        x = self.cbr(conv5)
        x = self.conv_last(x)

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
        else:
            x = nn.functional.log_softmax(x, dim=1)

        return x


# pyramid pooling, bilinear upsample
class PPMBilinear(nn.Module):
    def __init__(self, num_class=150, fc_dim=4096,
                 use_softmax=False, pool_scales=(1, 2, 3, 6)):
        super(PPMBilinear, self).__init__()
        self.use_softmax = use_softmax

        self.ppm = []
        for scale in pool_scales:
            self.ppm.append(nn.Sequential(
                nn.AdaptiveAvgPool2d(scale),
                nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(512),
                nn.ReLU(inplace=True)
            ))
        self.ppm = nn.ModuleList(self.ppm)

        self.conv_last = nn.Sequential(
            nn.Conv2d(fc_dim+len(pool_scales)*512, 512,
                      kernel_size=3, padding=1, bias=False, groups=cfg.RESNETS.NUM_GROUPS),
            SynchronizedBatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Dropout2d(0.1),
            nn.Conv2d(512, num_class, kernel_size=1)
        )


    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale in self.ppm:
            ppm_out.append(nn.functional.upsample(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False))
        ppm_out = torch.cat(ppm_out, 1)

        x = self.conv_last(ppm_out)


        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)

        return x


# pyramid pooling, bilinear upsample
class PPMBilinearDeepsup(nn.Module):
    def __init__(self, num_class=150, fc_dim=4096,
                 use_softmax=False, pool_scales=(1, 2, 3, 6)):
        super(PPMBilinearDeepsup, self).__init__()
        self.use_softmax = use_softmax

        self.ppm = []
        for scale in pool_scales:
            self.ppm.append(nn.Sequential(
                nn.AdaptiveAvgPool2d(scale),
                nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(512),
                nn.ReLU(inplace=True)
            ))
        self.ppm = nn.ModuleList(self.ppm)
        self.cbr_deepsup = conv3x3_bn_relu(fc_dim // 2, fc_dim // 4, 1)

        self.conv_last = nn.Sequential(
            nn.Conv2d(fc_dim+len(pool_scales)*512, 512,
                      kernel_size=3, padding=1, bias=False, groups=cfg.RESNETS.NUM_GROUPS),
            SynchronizedBatchNorm2d(512),
            nn.ReLU(inplace=True),
            nn.Dropout2d(0.1),
            nn.Conv2d(512, num_class, kernel_size=1)
        )
        self.conv_last_deepsup = nn.Conv2d(fc_dim // 4, num_class, 1, 1, 0)
        self.dropout_deepsup = nn.Dropout2d(0.1)


    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale in self.ppm:
            ppm_out.append(nn.functional.upsample(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False))
        ppm_out = torch.cat(ppm_out, 1)

        x = self.conv_last(ppm_out)

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        # deep sup
        conv4 = conv_out[-2]
        _ = self.cbr_deepsup(conv4)
        _ = self.dropout_deepsup(_)
        _ = self.conv_last_deepsup(_)


        return (x, _)



# upernet
class UPerNet(nn.Module):
    def __init__(self, num_class=150, fc_dim=4096,
                 use_softmax=False, pool_scales=(1, 2, 3, 6),
                 fpn_inplanes=(256,512,1024,2048), fpn_dim=256):
        super(UPerNet, self).__init__()
        self.use_softmax = use_softmax

        # PPM Module
        self.ppm_pooling = []
        self.ppm_conv = []

        for scale in pool_scales:
            self.ppm_pooling.append(nn.AdaptiveAvgPool2d(scale))
            self.ppm_conv.append(nn.Sequential(
                nn.Conv2d(fc_dim, 512, kernel_size=1,  bias=False),
                SynchronizedBatchNorm2d(512),
                nn.ReLU(inplace=True)
            ))
        self.ppm_pooling = nn.ModuleList(self.ppm_pooling)
        self.ppm_conv = nn.ModuleList(self.ppm_conv)
        if 'deform_conv' in cfg.MODEL.CONV_BODY and False:
            self.ppm_last_conv = conv3x3_bn_relu(fc_dim + len(pool_scales)*512, fpn_dim, 1)
            self.ppm_last_conv_drop=nn.Dropout2d(0.1)
        else:
            self.ppm_last_conv = conv3x3_bn_relu(fc_dim + len(pool_scales)*512, fpn_dim, 1)
        # FPN Module
        self.fpn_in = []
        for fpn_inplane in fpn_inplanes[:-1]: # skip the top layer
            self.fpn_in.append(nn.Sequential(
                nn.Conv2d(fpn_inplane, fpn_dim, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(fpn_dim),
                nn.ReLU(inplace=True)
            ))
        self.fpn_in = nn.ModuleList(self.fpn_in)

        self.fpn_out = []
        for i in range(len(fpn_inplanes) - 1): # skip the top layer
            if not cfg.SEM.GCN_ON:
                self.fpn_out.append(nn.Sequential(
                    conv3x3_bn_relu(fpn_dim, fpn_dim, 1),
                ))
            else:
                self.fpn_out.append(nn.Sequential(
                    GCN(fpn_dim, fpn_dim, kernel=15)
                ))
        self.fpn_out = nn.ModuleList(self.fpn_out)
        if 'deform_conv' in cfg.MODEL.CONV_BODY and False:
            self.conv_last = nn.Sequential(
                conv3x3_bn_relu(len(fpn_inplanes) * fpn_dim, fpn_dim, 1),
                nn.Dropout2d(0.9),
                nn.Conv2d(fpn_dim, num_class, kernel_size=1)
            )
        else:
            self.conv_last = nn.Sequential(
                conv3x3_bn_relu(len(fpn_inplanes) * fpn_dim, fpn_dim, 1),
                nn.Conv2d(fpn_dim, num_class, kernel_size=1)
            )

    def forward(self, conv_out, segSize=None):
        if (len(conv_out)>4):
            conv_out.pop(0)
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale, pool_conv in zip(self.ppm_pooling, self.ppm_conv):
            ppm_out.append(pool_conv(nn.functional.upsample(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False)))
        ppm_out = torch.cat(ppm_out, 1)
        f = self.ppm_last_conv(ppm_out)
        if 'deform_conv' in cfg.MODEL.CONV_BODY and False:
            f = self.ppm_last_conv_drop(f)
        fpn_feature_list = [f]
        for i in reversed(range(len(conv_out) - 1)):
            conv_x = conv_out[i]
            conv_x = self.fpn_in[i](conv_x) # lateral branch

            f = nn.functional.upsample(
                f, size=conv_x.size()[2:], mode='bilinear', align_corners=False) # top-down branch
            f = conv_x + f

            fpn_feature_list.append(self.fpn_out[i](f))

        fpn_feature_list.reverse() # [P2 - P5]
        output_size = fpn_feature_list[0].size()[2:]
        fusion_list = [fpn_feature_list[0]]
        for i in range(1, len(fpn_feature_list)):
            fusion_list.append(nn.functional.upsample(
                fpn_feature_list[i],
                output_size,
                mode='bilinear', align_corners=False))
        fusion_out = torch.cat(fusion_list, 1)
        x = self.conv_last(fusion_out)

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        return x


# deeplab
class DeepLabNet(nn.Module):
    def __init__(self, num_class=150, fc_dim=4096,
                 use_softmax=False, pool_scales=(1, 6, 12, 18),
                 fpn_inplanes=(256,2048), fpn_dim=256, low_dim=32):
        super(DeepLabNet, self).__init__()
        self.use_softmax = use_softmax
        # PPM Module
        self.ppm_pooling = []
        self.ppm_conv = []
        #add global pooling
        if cfg.SEM.USE_GP:
            self.ppm_pooling.append(nn.AdaptiveAvgPool2d(1))
            self.ppm_conv.append(nn.Sequential(
                nn.Conv2d(fc_dim, 256, kernel_size=1,  bias=False),
                SynchronizedBatchNorm2d(256),
                nn.ReLU(inplace=True)
            ))
        #add dilation conv
        for ids, scale in enumerate(pool_scales):
            if ids == 0:
                self.ppm_conv.append(nn.Sequential(
                    nn.Conv2d(fc_dim, 256, kernel_size=1,  groups=cfg.RESNETS.NUM_GROUPS, bias=False),
                    SynchronizedBatchNorm2d(256),
                    nn.ReLU(inplace=True)
                ))
            else:
                self.ppm_conv.append(nn.Sequential(
                    nn.Conv2d(fc_dim, 256, kernel_size=3, dilation=scale, padding=scale, groups=cfg.RESNETS.NUM_GROUPS, bias=False),
                    SynchronizedBatchNorm2d(256),
                    nn.ReLU(inplace=True)
                ))
        if cfg.SEM.GCN_ON:
            self.ppm_conv.append(nn.Sequential(
                GCN(fc_dim//2, 256, kernel=25),
                GCN(256, 256, kernel=25)))
        self.ppm_conv = nn.ModuleList(self.ppm_conv)
        self.ppm_pooling = nn.ModuleList(self.ppm_pooling)
        self.ppm_last_conv = nn.Sequential(
                conv3x3_bn_relu((len(pool_scales)+cfg.SEM.USE_GP+int(cfg.SEM.GCN_ON))*256, fpn_dim, 1),
                nn.Dropout(0.9))

        # FPN Module
        self.fpn_in = nn.Sequential(
                nn.Conv2d(fpn_inplanes[0], low_dim, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(low_dim),
                nn.ReLU(inplace=True)
            )
        self.fpn_out = conv3x3_bn_relu(low_dim+fpn_dim, fpn_dim, 1)
        self.conv_last = nn.Sequential(
            conv3x3_bn_relu(fpn_dim, fpn_dim, 1),
            nn.Conv2d(fpn_dim, num_class, kernel_size=1)
        )

    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = []
        #global pooling
        if cfg.SEM.USE_GP:
            ppm_out.append(self.ppm_conv[0](nn.functional.upsample(
                self.ppm_pooling[0](conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False)))
        #dilation conv
        for  pool_conv in self.ppm_conv[cfg.SEM.USE_GP: cfg.SEM.USE_GP+4]:
            ppm_out.append(pool_conv(conv5))
        if cfg.SEM.GCN_ON:
            ppm_out.append(self.ppm_conv[-1](conv_out[-2]))
        ppm_out = torch.cat(ppm_out, 1)
        f = self.ppm_last_conv(ppm_out)
        conv_x = conv_out[1]
        conv_x = self.fpn_in(conv_x) # lateral branch
        f = nn.functional.upsample(
            f, size=conv_x.size()[2:], mode='bilinear', align_corners=False) # top-down branch
        f = torch.cat([conv_x, f], dim=1)
        fpn_features = self.fpn_out(f)
        x = self.conv_last(fpn_features)
        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        return x

class Deeplab(nn.Module):
    def __init__(self, low_level_inplanes):
        super(Deeplab, self).__init__()

        #aspp
        self.aspp = build_aspp('', 16, BatchNorm)
        #decoder
        self.conv1 = nn.Conv2d(low_level_inplanes, 48, 1, bias=False, groups=cfg.RESNETS.NUM_GROUPS)
        self.bn1 = BatchNorm(48)
        self.relu = nn.ReLU()
        self.last_conv = nn.Sequential(nn.Conv2d(304, 256, kernel_size=3, stride=1, padding=1, bias=False, groups=cfg.RESNETS.NUM_GROUPS),
                                       BatchNorm(256),
                                       nn.ReLU(),
                                       nn.Dropout(0.5),
                                       nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1, bias=False, groups=cfg.RESNETS.NUM_GROUPS),
                                       BatchNorm(256),
                                       nn.ReLU(),
                                       nn.Dropout(0.1),
                                       nn.Conv2d(256, cfg.MODEL.NUM_CLASSES, kernel_size=1, stride=1))
        self._init_weight()


    def forward(self, conv_out):
        low_level_feat = conv_out[1]
        x = self.aspp(conv_out[-1])
        low_level_feat = self.conv1(low_level_feat)
        low_level_feat = self.bn1(low_level_feat)
        low_level_feat = self.relu(low_level_feat)

        x = F.interpolate(x, size=low_level_feat.size()[2:], mode='bilinear', align_corners=True)
        x = torch.cat((x, low_level_feat), dim=1)
        x = self.last_conv(x)
        

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                spn_x, size=cfg.SEM.INPUT_SIZE, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        return x

    def _init_weight(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                torch.nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, SynchronizedBatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

class CspnUPerNet(nn.Module):
    def __init__(self, num_class=150, fc_dim=4096,
                 use_softmax=False, pool_scales=(1, 2, 3, 6),
                 fpn_inplanes=(256,512,1024,2048), fpn_dim=256):
        super(CspnUPerNet, self).__init__()
        self.use_softmax = use_softmax

        # PPM Module
        self.ppm_pooling = []
        self.ppm_conv = []

        for scale in pool_scales:
            self.ppm_pooling.append(nn.AdaptiveAvgPool2d(scale))
            self.ppm_conv.append(nn.Sequential(
                nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(512),
                nn.ReLU(inplace=True)
            ))
        self.ppm_pooling = nn.ModuleList(self.ppm_pooling)
        self.ppm_conv = nn.ModuleList(self.ppm_conv)
        self.ppm_last_conv = nn.Sequential(
                conv3x3_bn_relu(fc_dim + len(pool_scales)*512, fpn_dim, 1),
                nn.Dropout2d(0.1))

        # FPN Module
        self.fpn_in = []
        for fpn_inplane in fpn_inplanes[:-1]: # skip the top layer
            self.fpn_in.append(nn.Sequential(
                nn.Conv2d(fpn_inplane, fpn_dim, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(fpn_dim),
                nn.ReLU(inplace=True)
            ))
        self.fpn_in = nn.ModuleList(self.fpn_in)

        self.fpn_out = []
        for i in range(len(fpn_inplanes) - 1): # skip the top layer
            self.fpn_out.append(nn.Sequential(
                conv3x3_bn_relu(fpn_dim, fpn_dim, 1),
            ))
        self.fpn_out = nn.ModuleList(self.fpn_out)

        self.fpn_out_last = nn.Sequential(
            conv3x3_bn_relu(len(fpn_inplanes) * fpn_dim, fpn_dim, 1),
            nn.Dropout2d(0.9),
        )
        self.conv_last = nn.Conv2d(fpn_dim, num_class, kernel_size=1)

        self.cspn_last = nn.Sequential(
            nn.Conv2d(fpn_dim, 8*cfg.MODEL.NUM_CLASSES, 3, padding=1, groups=cfg.RESNETS.NUM_GROUPS),
            SynchronizedBatchNorm2d(8*cfg.MODEL.NUM_CLASSES),
        )
        self.cspn_net = Affinity_Propagate()


    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale, pool_conv in zip(self.ppm_pooling, self.ppm_conv):
            ppm_out.append(pool_conv(nn.functional.upsample(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False)))
        ppm_out = torch.cat(ppm_out, 1)
        f = self.ppm_last_conv(ppm_out)

        fpn_feature_list = [f]
        for i in reversed(range(len(conv_out) - 1)):
            conv_x = conv_out[i]
            conv_x = self.fpn_in[i](conv_x) # lateral branch

            f = nn.functional.upsample(
                f, size=conv_x.size()[2:], mode='bilinear', align_corners=False) # top-down branch
            f = conv_x + f

            fpn_feature_list.append(self.fpn_out[i](f))

        fpn_feature_list.reverse() # [P2 - P5]
        output_size = fpn_feature_list[0].size()[2:]
        fusion_list = [fpn_feature_list[0]]
        for i in range(1, len(fpn_feature_list)):
            fusion_list.append(nn.functional.upsample(
                fpn_feature_list[i],
                output_size,
                mode='bilinear', align_corners=False))
        fusion_out = torch.cat(fusion_list, 1)
        fusion_out = self.fusion_out_last(fusion_out)
        fusion_out = nn.functional.upsample(fusion_out, scale_factor=2, 
                mode='bilinear', align_corners=False)

        x = self.conv_last(fusion_out)
        #x= nn.functional.upsample(x, scale_factor=2, mode='bilinear', align_corners=False)
        x = nn.functional.softmax(x, dim=1)

        guidance = self.cspn_last(fusion_out)
        #guidance= nn.functional.upsample(guidance, scale_factor=2, mode='bilinear', align_corners=False)

        b, c, h , w = guidance.shape
        guidance = guidance.view(b, 8, -1, h, w)
        x = self.cspn_net(guidance, x)

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                x, size=segSize, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        return x

# upernet+spn
class SpnUPerNet(nn.Module):
    def __init__(self, num_class=150, fc_dim=4096,
                 use_softmax=False, pool_scales=(1, 2, 3, 6),
                 fpn_inplanes=(256,512,1024,2048), fpn_dim=256):
        super(SpnUPerNet, self).__init__()
        self.use_softmax = use_softmax

        # PPM Module
        self.ppm_pooling = []
        self.ppm_conv = []

        for scale in pool_scales:
            self.ppm_pooling.append(nn.AdaptiveAvgPool2d(scale))
            self.ppm_conv.append(nn.Sequential(
                nn.Conv2d(fc_dim, 512, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(512),
                nn.ReLU(inplace=True)
            ))
        self.ppm_pooling = nn.ModuleList(self.ppm_pooling)
        self.ppm_conv = nn.ModuleList(self.ppm_conv)
        self.ppm_last_conv = conv3x3_bn_relu(fc_dim + len(pool_scales)*512, fpn_dim, 1)

        # FPN Module
        self.fpn_in = []
        for fpn_inplane in fpn_inplanes[:-1]: # skip the top layer
            self.fpn_in.append(nn.Sequential(
                nn.Conv2d(fpn_inplane, fpn_dim, kernel_size=1, bias=False),
                SynchronizedBatchNorm2d(fpn_dim),
                nn.ReLU(inplace=True)
            ))
        self.fpn_in = nn.ModuleList(self.fpn_in)

        self.fpn_out = []
        for i in range(len(fpn_inplanes) - 1): # skip the top layer
            self.fpn_out.append(nn.Sequential(
                conv3x3_bn_relu(fpn_dim, fpn_dim, 1),
            ))
        self.fpn_out = nn.ModuleList(self.fpn_out)

        self.conv_last = nn.Sequential(
            conv3x3_bn_relu(len(fpn_inplanes) * fpn_dim, fpn_dim, 1),
            nn.Conv2d(fpn_dim, num_class, kernel_size=1)
        )
        self.spn_net = SPN()


    def forward(self, conv_out, segSize=None):
        conv5 = conv_out[-1]

        input_size = conv5.size()
        ppm_out = [conv5]
        for pool_scale, pool_conv in zip(self.ppm_pooling, self.ppm_conv):
            ppm_out.append(pool_conv(nn.functional.upsample(
                pool_scale(conv5),
                (input_size[2], input_size[3]),
                mode='bilinear', align_corners=False)))
        ppm_out = torch.cat(ppm_out, 1)
        f = self.ppm_last_conv(ppm_out)

        fpn_feature_list = [f]
        for i in reversed(range(len(conv_out) - 1)):
            conv_x = conv_out[i]
            conv_x = self.fpn_in[i](conv_x) # lateral branch

            f = nn.functional.upsample(
                f, size=conv_x.size()[2:], mode='bilinear', align_corners=False) # top-down branch
            f = conv_x + f

            fpn_feature_list.append(self.fpn_out[i](f))

        fpn_feature_list.reverse() # [P2 - P5]
        output_size = fpn_feature_list[0].size()[2:]
        fusion_list = [fpn_feature_list[0]]
        for i in range(1, len(fpn_feature_list)):
            fusion_list.append(nn.functional.upsample(
                fpn_feature_list[i],
                output_size,
                mode='bilinear', align_corners=False))
        fusion_out = torch.cat(fusion_list, 1)
        x = self.conv_last(fusion_out)
        x = nn.functional.upsample(x, scale_factor=4, mode='bilinear', align_corners=False)

        guidance= nn.functional.upsample(fusion_out, scale_factor=2, mode='bilinear', align_corners=False)
        spn_x = self.spn_net(x, guidance)

        if self.use_softmax or segSize is not None:  # is True during inference
            x = nn.functional.upsample(
                spn_x, size=cfg.SEM.INPUT_SIZE, mode='bilinear', align_corners=False)
            x = nn.functional.softmax(x, dim=1)
            return x

        return spn_x, x
