<gh_stars>0
import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F

class fcn32s_partial(nn.Module):

    def __init__(self):

        super(fcn32s_partial, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3,out_channels=64,kernel_size=3,stride=1,padding=100),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64,out_channels=64,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128,out_channels=128,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=128,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(in_channels=256,out_channels=512,kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels= 512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2,ceil_mode=True),
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(in_channels=512,out_channels=512,kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512,out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512,out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )

    def forward(self,x):

        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        return x

class fcn32s(nn.Module):
    def __init__(self,class_number=21,transition=True):
        super(fcn32s, self).__init__()

        self.features = fcn32s_partial()
        self.transition = transition
        if self.transition:

            self.transition_layer = nn.Sequential(
                nn.Conv2d(in_channels=512*2,out_channels=512,kernel_size=3,padding=1),
                nn.ReLU(inplace=True)
            )
            self.fc6 = nn.Sequential(
               nn.Conv2d(in_channels=512,out_channels=4096,kernel_size=7),
               nn.ReLU(inplace=True),
               nn.Dropout2d(),
            )
        else:
            self.fc6 = nn.Sequential(
                nn.Conv2d(in_channels=512 * 2, out_channels=4096, kernel_size=7),
                nn.ReLU(inplace=True),
                nn.Dropout2d(),
            )
        self.fc7 = nn.Sequential(
            nn.Conv2d(in_channels=4096,out_channels=4096,kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(),
        )
        self.score = nn.Conv2d(in_channels=4096,out_channels=class_number,kernel_size=1)
        self.upscore = nn.ConvTranspose2d(in_channels=class_number,out_channels=class_number,kernel_size=64,stride=32,bias=False)

    def forward(self,t0,t1):

        feat_t0 = self.features(t0)
        feat_t1 = self.features(t1)

        if self.transition:
            x = self.transition_layer(torch.cat((feat_t0,feat_t1),dim=1))
            x = self.fc6(x)
            x = self.fc7(x)
        else:
            x = self.fc6(torch.cat((feat_t0,feat_t1),dim=1))
            x = self.fc7(x)
        score = self.score(x)
        upscore = self.upscore(score)
        upscore = upscore[:, :,19:19+t0.size()[2], 19:19+t0.size()[3]].contiguous()
        return upscore

    def init_parameters(self,pretrain_vgg16):

        ##### init parameter using pretrain vgg16 model ###########
        conv_blocks = [self.features.conv1,
                       self.features.conv2,
                       self.features.conv3,
                       self.features.conv4,
                       self.features.conv5]

        ranges = [[0, 4], [5, 9], [10, 16], [17, 23], [24, 29]]
        features = list(pretrain_vgg16.features.children())

        for idx, conv_block in enumerate(conv_blocks):
            for l1, l2 in zip(features[ranges[idx][0]:ranges[idx][1]], conv_block):
                if isinstance(l1, nn.Conv2d) and isinstance(l2, nn.Conv2d):
                    # print idx, l1, l2
                    assert l1.weight.size() == l2.weight.size()
                    assert l1.bias.size() == l2.bias.size()
                    l2.weight.data = l1.weight.data
                    l2.bias.data = l1.bias.data

        ####### init fc parameters (transplant) ##############
        if self.transition:
            self.transition_layer[0].weight.data.normal_(0, 0.01)
            self.transition_layer[0].bias.data.fill_(0)

            self.fc6[0].weight.data = pretrain_vgg16.classifier[0].weight.data.view(self.fc6[0].weight.size())
            self.fc6[0].bias.data = pretrain_vgg16.classifier[0].bias.data.view(self.fc6[0].bias.size())

        else:
            self.fc6[0].weight.data = torch.cat((pretrain_vgg16.classifier[0].weight.data,
                                                 pretrain_vgg16.classifier[0].weight.data), dim=1)

            self.fc6[0].bias.data = pretrain_vgg16.classifier[0].bias.data.view(self.fc6[0].bias.size())

        self.fc7[0].weight.data = pretrain_vgg16.classifier[3].weight.data.view(self.fc7[0].weight.size())
        self.fc7[0].bias.data = pretrain_vgg16.classifier[3].bias.data.view(self.fc7[0].bias.size())

        ###### random init socore layer parameters ###########
        assert  self.upscore.kernel_size[0] == self.upscore.kernel_size[1]
        initial_weight = get_upsampling_weight(self.upscore.in_channels, self.upscore.out_channels, self.upscore.kernel_size[0])
        self.upscore.weight.data.copy_(initial_weight)

# https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/surgery.py
def get_upsampling_weight(in_channels, out_channels, kernel_size):
    """Make a 2D bilinear kernel suitable for upsampling"""
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
                      dtype=np.float64)
    weight[range(in_channels), range(out_channels), :, :] = filt
    return torch.from_numpy(weight).float()

import torch
import torch.nn as nn
import numpy as np
import torch.nn.functional as F

class fcn32s_partial(nn.Module):

    def __init__(self):

        super(fcn32s_partial, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3,out_channels=64,kernel_size=3,stride=1,padding=100),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=64,out_channels=64,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=64,out_channels=128,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=128,out_channels=128,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=128,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256,out_channels=256,kernel_size=3,stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(in_channels=256,out_channels=512,kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512, out_channels= 512, kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2,ceil_mode=True),
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(in_channels=512,out_channels=512,kernel_size=3, stride=1,padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512,out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=512,out_channels=512, kernel_size=3, stride=1, padding=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2,stride=2,ceil_mode=True),
        )

    def forward(self,x):

        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        return x

class fcn32s(nn.Module):
    def __init__(self,class_number=21,transition=True):
        super(fcn32s, self).__init__()

        self.features = fcn32s_partial()
        self.transition = transition
        if self.transition:

            self.transition_layer = nn.Sequential(
                nn.Conv2d(in_channels=512*2,out_channels=512,kernel_size=3,padding=1),
                nn.ReLU(inplace=True)
            )
            self.fc6 = nn.Sequential(
               nn.Conv2d(in_channels=512,out_channels=4096,kernel_size=7),
               nn.ReLU(inplace=True),
               nn.Dropout2d(),
            )
        else:
            self.fc6 = nn.Sequential(
                nn.Conv2d(in_channels=512 * 2, out_channels=4096, kernel_size=7),
                nn.ReLU(inplace=True),
                nn.Dropout2d(),
            )
        self.fc7 = nn.Sequential(
            nn.Conv2d(in_channels=4096,out_channels=4096,kernel_size=1),
            nn.ReLU(inplace=True),
            nn.Dropout2d(),
        )
        self.score = nn.Conv2d(in_channels=4096,out_channels=class_number,kernel_size=1)
        self.upscore = nn.ConvTranspose2d(in_channels=class_number,out_channels=class_number,kernel_size=64,stride=32,bias=False)

    def forward(self,t0,t1):

        feat_t0 = self.features(t0)
        feat_t1 = self.features(t1)

        if self.transition:
            x = self.transition_layer(torch.cat((feat_t0,feat_t1),dim=1))
            x = self.fc6(x)
            x = self.fc7(x)
        else:
            x = self.fc6(torch.cat((feat_t0,feat_t1),dim=1))
            x = self.fc7(x)
        score = self.score(x)
        upscore = self.upscore(score)
        upscore = upscore[:, :,19:19+t0.size()[2], 19:19+t0.size()[3]].contiguous()
        return upscore

    def init_parameters(self,pretrain_vgg16):

        ##### init parameter using pretrain vgg16 model ###########
        conv_blocks = [self.features.conv1,
                       self.features.conv2,
                       self.features.conv3,
                       self.features.conv4,
                       self.features.conv5]

        ranges = [[0, 4], [5, 9], [10, 16], [17, 23], [24, 29]]
        features = list(pretrain_vgg16.features.children())

        for idx, conv_block in enumerate(conv_blocks):
            for l1, l2 in zip(features[ranges[idx][0]:ranges[idx][1]], conv_block):
                if isinstance(l1, nn.Conv2d) and isinstance(l2, nn.Conv2d):
                    # print idx, l1, l2
                    assert l1.weight.size() == l2.weight.size()
                    assert l1.bias.size() == l2.bias.size()
                    l2.weight.data = l1.weight.data
                    l2.bias.data = l1.bias.data

        ####### init fc parameters (transplant) ##############
        if self.transition:
            self.transition_layer[0].weight.data.normal_(0, 0.01)
            self.transition_layer[0].bias.data.fill_(0)

            self.fc6[0].weight.data = pretrain_vgg16.classifier[0].weight.data.view(self.fc6[0].weight.size())
            self.fc6[0].bias.data = pretrain_vgg16.classifier[0].bias.data.view(self.fc6[0].bias.size())

        else:
            self.fc6[0].weight.data = torch.cat((pretrain_vgg16.classifier[0].weight.data,
                                                 pretrain_vgg16.classifier[0].weight.data), dim=1)

            self.fc6[0].bias.data = pretrain_vgg16.classifier[0].bias.data.view(self.fc6[0].bias.size())

        self.fc7[0].weight.data = pretrain_vgg16.classifier[3].weight.data.view(self.fc7[0].weight.size())
        self.fc7[0].bias.data = pretrain_vgg16.classifier[3].bias.data.view(self.fc7[0].bias.size())

        ###### random init socore layer parameters ###########
        assert  self.upscore.kernel_size[0] == self.upscore.kernel_size[1]
        initial_weight = get_upsampling_weight(self.upscore.in_channels, self.upscore.out_channels, self.upscore.kernel_size[0])
        self.upscore.weight.data.copy_(initial_weight)

# https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/surgery.py
def get_upsampling_weight(in_channels, out_channels, kernel_size):
    """Make a 2D bilinear kernel suitable for upsampling"""
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * \
           (1 - abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
                      dtype=np.float64)
    weight[range(in_channels), range(out_channels), :, :] = filt
    return torch.from_numpy(weight).float()
##### init parameter using pretrain vgg16 model ########### # print idx, l1, l2 ####### init fc parameters (transplant) ############## ###### random init socore layer parameters ########### # https://github.com/shelhamer/fcn.berkeleyvision.org/blob/master/surgery.py Make a 2D bilinear kernel suitable for upsampling