trafficDetectionHandover/models/bisenet.py

"""Bilateral Segmentation Network"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
# from core.models.base_models.resnet import resnet18,resnet50
from torchvision import models
# from core.nn import _ConvBNReLU

# __all__ = ['BiSeNet', 'get_bisenet', 'get_bisenet_resnet18_citys']

class _ConvBNReLU(nn.Module):
    def __init__(self,in_channels,out_channels, k, s, p, norm_layer=None):
        super(_ConvBNReLU, self).__init__()
        self.conv =nn.Conv2d(in_channels, out_channels, kernel_size=k, stride=s, padding=p)
        self.bn = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace = True)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)

        return x
class BiSeNet(nn.Module):
    def __init__(self, nclass, backbone='resnet18', aux=False, jpu=False, pretrained_base=True, **kwargs):
        super(BiSeNet, self).__init__()
        self.aux = aux
        self.spatial_path = SpatialPath(3, 128, **kwargs)
        self.context_path = ContextPath(backbone, pretrained_base, **kwargs)
        self.ffm = FeatureFusion(256, 256, 4, **kwargs)
        self.head = _BiSeHead(256, 64, nclass, **kwargs)
        if aux:
            self.auxlayer1 = _BiSeHead(128, 256, nclass, **kwargs)
            self.auxlayer2 = _BiSeHead(128, 256, nclass, **kwargs)

        self.__setattr__('exclusive',
                         ['spatial_path', 'context_path', 'ffm', 'head', 'auxlayer1', 'auxlayer2'] if aux else [
                             'spatial_path', 'context_path', 'ffm', 'head'])

    def forward(self, x,outsize=None,test_flag=False):
        size = x.size()[2:]
        spatial_out = self.spatial_path(x)
        context_out = self.context_path(x)
        fusion_out = self.ffm(spatial_out, context_out[-1])
        outputs = []
        x = self.head(fusion_out)
        x = F.interpolate(x, size, mode='bilinear', align_corners=True)

        if outsize:
            print('######using   torch resize#######',outsize)
            x = F.interpolate(x, outsize, mode='bilinear', align_corners=True)
        outputs.append(x)

        if self.aux:
            auxout1 = self.auxlayer1(context_out[0])
            auxout1 = F.interpolate(auxout1, size, mode='bilinear', align_corners=True)
            outputs.append(auxout1)
            auxout2 = self.auxlayer2(context_out[1])
            auxout2 = F.interpolate(auxout2, size, mode='bilinear', align_corners=True)
            outputs.append(auxout2)
        if test_flag:
            outputs = [torch.argmax(outputx, axis=1) for outputx in outputs]
        #return tuple(outputs)
        return outputs[0]

class BiSeNet_MultiOutput(nn.Module):
    def __init__(self, nclass, backbone='resnet18', aux=False, jpu=False, pretrained_base=True, **kwargs):
        super(BiSeNet_MultiOutput, self).__init__()
        self.aux = aux
        self.spatial_path = SpatialPath(3, 128, **kwargs)
        self.context_path = ContextPath(backbone, pretrained_base, **kwargs)
        self.ffm = FeatureFusion(256, 256, 4, **kwargs)
        assert isinstance(nclass, list)
        self.outCnt = len(nclass)
        for ii, nclassii in enumerate(nclass):
            setattr(self, 'head%d'%(ii), _BiSeHead(256, 64, nclassii, **kwargs))

        if aux:
            self.auxlayer1 = _BiSeHead(128, 256, nclass, **kwargs)
            self.auxlayer2 = _BiSeHead(128, 256, nclass, **kwargs)

        self.__setattr__('exclusive',
                         ['spatial_path', 'context_path', 'ffm', 'head', 'auxlayer1', 'auxlayer2'] if aux else [
                             'spatial_path', 'context_path', 'ffm', 'head'])

    def forward(self, x, outsize=None, test_flag=False, smooth_kernel=0):
        size = x.size()[2:]
        spatial_out = self.spatial_path(x)
        context_out = self.context_path(x)
        fusion_out = self.ffm(spatial_out, context_out[-1])
        outputs = []
        for ii in range(self.outCnt):
            x = getattr(self, 'head%d'%(ii))(fusion_out)
            x = F.interpolate(x, size, mode='bilinear', align_corners=True)
            outputs.append(x)
      
        if self.aux:
            auxout1 = self.auxlayer1(context_out[0])
            auxout1 = F.interpolate(auxout1, size, mode='bilinear', align_corners=True)
            outputs.append(auxout1)
            auxout2 = self.auxlayer2(context_out[1])
            auxout2 = F.interpolate(auxout2, size, mode='bilinear', align_corners=True)
            outputs.append(auxout2)
        if test_flag:
            outputs = [torch.argmax(outputx ,axis=1)  for outputx in outputs]  
        if smooth_kernel>0:
            gaussian_kernel = torch.from_numpy(np.ones((1,1,smooth_kernel,smooth_kernel))  )
            
            pad = int((smooth_kernel - 1)/2)
            if not gaussian_kernel.is_cuda:
                gaussian_kernel = gaussian_kernel.to(x.device)
                #print(gaussian_kernel.dtype,gaussian_kernel,outputs[0].dtype)           
            outputs = [x.unsqueeze(1).double() for x in outputs]
            outputs = [torch.conv2d(x, gaussian_kernel, padding=pad) for x in outputs]
            outputs = [x.squeeze(1).long() for x in outputs]
        #return tuple(outputs)
        return outputs
        
class _BiSeHead(nn.Module):
    def __init__(self, in_channels, inter_channels, nclass, norm_layer=nn.BatchNorm2d, **kwargs):
        super(_BiSeHead, self).__init__()
        self.block = nn.Sequential(
            _ConvBNReLU(in_channels, inter_channels, 3, 1, 1, norm_layer=norm_layer),
            nn.Dropout(0.1),
            nn.Conv2d(inter_channels, nclass, 1)
        )

    def forward(self, x):
        x = self.block(x)
        return x


class SpatialPath(nn.Module):
    """Spatial path"""

    def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d, **kwargs):
        super(SpatialPath, self).__init__()
        inter_channels = 64
        self.conv7x7 = _ConvBNReLU(in_channels, inter_channels, 7, 2, 3, norm_layer=norm_layer)
        self.conv3x3_1 = _ConvBNReLU(inter_channels, inter_channels, 3, 2, 1, norm_layer=norm_layer)
        self.conv3x3_2 = _ConvBNReLU(inter_channels, inter_channels, 3, 2, 1, norm_layer=norm_layer)
        self.conv1x1 = _ConvBNReLU(inter_channels, out_channels, 1, 1, 0, norm_layer=norm_layer)

    def forward(self, x):
        x = self.conv7x7(x)
        x = self.conv3x3_1(x)
        x = self.conv3x3_2(x)
        x = self.conv1x1(x)

        return x


class _GlobalAvgPooling(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer, **kwargs):
        super(_GlobalAvgPooling, self).__init__()
        self.gap = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Conv2d(in_channels, out_channels, 1, bias=False),
            norm_layer(out_channels),
            nn.ReLU(True)
        )

    def forward(self, x):
        size = x.size()[2:]
        pool = self.gap(x)
        out = F.interpolate(pool, size, mode='bilinear', align_corners=True)
        return out


class AttentionRefinmentModule(nn.Module):
    def __init__(self, in_channels, out_channels, norm_layer=nn.BatchNorm2d, **kwargs):
        super(AttentionRefinmentModule, self).__init__()
        self.conv3x3 = _ConvBNReLU(in_channels, out_channels, 3, 1, 1, norm_layer=norm_layer)
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            _ConvBNReLU(out_channels, out_channels, 1, 1, 0, norm_layer=norm_layer),
            nn.Sigmoid()
        )

    def forward(self, x):
        x = self.conv3x3(x)
        attention = self.channel_attention(x)
        x = x * attention
        return x


class ContextPath(nn.Module):
    def __init__(self, backbone='resnet18', pretrained_base=True, norm_layer=nn.BatchNorm2d, **kwargs):
        super(ContextPath, self).__init__()
        if backbone == 'resnet18':
            pretrained = models.resnet18(pretrained=pretrained_base, **kwargs)
        elif backbone=='resnet50':
            pretrained = models.resnet50(pretrained=pretrained_base, **kwargs)
        else:
            raise RuntimeError('unknown backbone: {}'.format(backbone))
        self.conv1 = pretrained.conv1
        self.bn1 = pretrained.bn1
        self.relu = pretrained.relu
        self.maxpool = pretrained.maxpool
        self.layer1 = pretrained.layer1
        self.layer2 = pretrained.layer2
        self.layer3 = pretrained.layer3
        self.layer4 = pretrained.layer4

        inter_channels = 128
        self.global_context = _GlobalAvgPooling(512, inter_channels, norm_layer)

        self.arms = nn.ModuleList(
            [AttentionRefinmentModule(512, inter_channels, norm_layer, **kwargs),
             AttentionRefinmentModule(256, inter_channels, norm_layer, **kwargs)]
        )
        self.refines = nn.ModuleList(
            [_ConvBNReLU(inter_channels, inter_channels, 3, 1, 1, norm_layer=norm_layer),
             _ConvBNReLU(inter_channels, inter_channels, 3, 1, 1, norm_layer=norm_layer)]
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)
        x = self.layer1(x)

        context_blocks = []
        context_blocks.append(x)
        x = self.layer2(x)
        context_blocks.append(x)
        c3 = self.layer3(x)
        context_blocks.append(c3)
        c4 = self.layer4(c3)
        context_blocks.append(c4)
        context_blocks.reverse()

        global_context = self.global_context(c4)
        last_feature = global_context
        context_outputs = []
        for i, (feature, arm, refine) in enumerate(zip(context_blocks[:2], self.arms, self.refines)):
            feature = arm(feature)
            feature += last_feature
            last_feature = F.interpolate(feature, size=context_blocks[i + 1].size()[2:],
                                         mode='bilinear', align_corners=True)
            last_feature = refine(last_feature)
            context_outputs.append(last_feature)

        return context_outputs


class FeatureFusion(nn.Module):
    def __init__(self, in_channels, out_channels, reduction=1, norm_layer=nn.BatchNorm2d, **kwargs):
        super(FeatureFusion, self).__init__()
        self.conv1x1 = _ConvBNReLU(in_channels, out_channels, 1, 1, 0, norm_layer=norm_layer, **kwargs)
        self.channel_attention = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            _ConvBNReLU(out_channels, out_channels // reduction, 1, 1, 0, norm_layer=norm_layer),
            _ConvBNReLU(out_channels // reduction, out_channels, 1, 1, 0, norm_layer=norm_layer),
            nn.Sigmoid()
        )

    def forward(self, x1, x2):
        fusion = torch.cat([x1, x2], dim=1)
        out = self.conv1x1(fusion)
        attention = self.channel_attention(out)
        out = out + out * attention
        return out


# def get_bisenet(dataset='citys', backbone='resnet18', pretrained=False, root='~/.torch/models',
#                 pretrained_base=True, **kwargs):
#     acronyms = {
#         'pascal_voc': 'pascal_voc',
#         'pascal_aug': 'pascal_aug',
#         'ade20k': 'ade',
#         'coco': 'coco',
#         'citys': 'citys',
#     }
#     from ..data.dataloader import datasets
#     model = BiSeNet(datasets[dataset].NUM_CLASS, backbone=backbone, pretrained_base=pretrained_base, **kwargs)
#     if pretrained:
#         from .model_store import get_model_file
#         device = torch.device(kwargs['local_rank'])
#         model.load_state_dict(torch.load(get_model_file('bisenet_%s_%s' % (backbone, acronyms[dataset]), root=root),
#                               map_location=device))
#     return model
#
#
# def get_bisenet_resnet18_citys(**kwargs):
#     return get_bisenet('citys', 'resnet18', **kwargs)


# if __name__ == '__main__':
#     # img = torch.randn(2, 3, 224, 224)
#     # model = BiSeNet(19, backbone='resnet18')
#     # print(model.exclusive)
#     input = torch.rand(2, 3, 224, 224)
#     model = BiSeNet(4, pretrained_base=True)
#     # target = torch.zeros(4, 512, 512).cuda()
#     # model.eval()
#     # print(model)
#     loss = model(input)
#     print(loss, loss.shape)
#
#     # from torchsummary import summary
#     #
#     # summary(model, (3, 224, 224))  # 打印表格，按顺序输出每层的输出形状和参数
#     import torch
#     from thop import profile
#     from torchsummary import summary
#
#     flop, params = profile(model, input_size=(1, 3, 512, 512))
#     print('flops:{:.3f}G\nparams:{:.3f}M'.format(flop / 1e9, params / 1e6))

if __name__ == '__main__':
    x = torch.rand(2, 3, 256, 256)

    # model = BiSeNet_MultiOutput(nclass=[2, 2])  # 原始
    # model = BiSeNet_MultiOutput(nclass=[3, 3])  # 改动
    model = BiSeNet_MultiOutput(nclass=[3, 3])  # 改动

    # print(model)
    out = model(x)
    print(out[0].size())
    # print()