Upload models/network.py

models/network.py

@@ -0,0 +1,352 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import init
+import torchvision
+import torch.nn.utils.spectral_norm as spectral_norm
+import math
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, inChannels, outChannels, convNum, normLayer=None):
+        super(ConvBlock, self).__init__()
+        self.inConv = nn.Sequential(
+            nn.Conv2d(inChannels, outChannels, kernel_size=3, padding=1),
+            nn.ReLU(inplace=True)
+        )
+        layers = []
+        for _ in range(convNum - 1):
+            layers.append(nn.Conv2d(outChannels, outChannels, kernel_size=3, padding=1))
+            layers.append(nn.ReLU(inplace=True))
+        if not (normLayer is None):
+            layers.append(normLayer(outChannels))
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.inConv(x)
+        x = self.conv(x)
+        return x
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels, normLayer=None):
+        super(ResidualBlock, self).__init__()
+        layers = []
+        layers.append(nn.Conv2d(channels, channels, kernel_size=3, padding=1))
+        layers.append(spectral_norm(nn.Conv2d(channels, channels, kernel_size=3, padding=1)))
+        if not (normLayer is None):
+            layers.append(normLayer(channels))
+        layers.append(nn.ReLU(inplace=True))
+        layers.append(nn.Conv2d(channels, channels, kernel_size=3, padding=1))
+        if not (normLayer is None):
+            layers.append(normLayer(channels))
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        residual = self.conv(x)
+        return F.relu(x + residual, inplace=True)
+
+
+class ResidualBlockSN(nn.Module):
+    def __init__(self, channels, normLayer=None):
+        super(ResidualBlockSN, self).__init__()
+        layers = []
+        layers.append(spectral_norm(nn.Conv2d(channels, channels, kernel_size=3, padding=1)))
+        layers.append(nn.LeakyReLU(0.2, True))
+        layers.append(spectral_norm(nn.Conv2d(channels, channels, kernel_size=3, padding=1)))
+        if not (normLayer is None):
+            layers.append(normLayer(channels))
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        residual = self.conv(x)
+        return F.leaky_relu(x + residual, 2e-1, inplace=True)
+
+
+class DownsampleBlock(nn.Module):
+    def __init__(self, inChannels, outChannels, convNum=2, normLayer=None):
+        super(DownsampleBlock, self).__init__()
+        layers = []
+        layers.append(nn.Conv2d(inChannels, outChannels, kernel_size=3, padding=1, stride=2))
+        layers.append(nn.ReLU(inplace=True))
+        for _ in range(convNum - 1):
+            layers.append(nn.Conv2d(outChannels, outChannels, kernel_size=3, padding=1))
+            layers.append(nn.ReLU(inplace=True))
+        if not (normLayer is None):
+            layers.append(normLayer(outChannels))
+        self.conv = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.conv(x)
+
+
+class UpsampleBlock(nn.Module):
+    def __init__(self, inChannels, outChannels, convNum=2, normLayer=None):
+        super(UpsampleBlock, self).__init__()
+        self.conv1 = nn.Conv2d(inChannels, outChannels, kernel_size=3, padding=1, stride=1)
+        self.combine = nn.Conv2d(2 * outChannels, outChannels, kernel_size=3, padding=1)
+        layers = []
+        for _ in range(convNum - 1):
+            layers.append(nn.Conv2d(outChannels, outChannels, kernel_size=3, padding=1))
+            layers.append(nn.ReLU(inplace=True))
+        if not (normLayer is None):
+            layers.append(normLayer(outChannels))
+        self.conv2 = nn.Sequential(*layers)
+
+    def forward(self, x, x0):
+        x = self.conv1(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = self.combine(torch.cat((x, x0), 1))
+        x = F.relu(x)
+        return self.conv2(x)
+
+
+class UpsampleBlockSN(nn.Module):
+    def __init__(self, inChannels, outChannels, convNum=2, normLayer=None):
+        super(UpsampleBlockSN, self).__init__()
+        self.conv1 = spectral_norm(nn.Conv2d(inChannels, outChannels, kernel_size=3, stride=1, padding=1))
+        self.shortcut = spectral_norm(nn.Conv2d(outChannels, outChannels, kernel_size=3, stride=1, padding=1))
+        layers = []
+        for _ in range(convNum - 1):
+            layers.append(spectral_norm(nn.Conv2d(outChannels, outChannels, kernel_size=3, padding=1)))
+            layers.append(nn.LeakyReLU(0.2, True))
+        if not (normLayer is None):
+            layers.append(normLayer(outChannels))
+        self.conv2 = nn.Sequential(*layers)
+
+    def forward(self, x, x0):
+        x = self.conv1(x)
+        x = F.interpolate(x, scale_factor=2, mode='nearest')
+        x = x + self.shortcut(x0)
+        x = F.leaky_relu(x, 2e-1)
+        return self.conv2(x)
+
+
+class HourGlass2(nn.Module):
+    def __init__(self, inChannel=3, outChannel=1, resNum=3, normLayer=None):
+        super(HourGlass2, self).__init__()
+        self.inConv = ConvBlock(inChannel, 64, convNum=2, normLayer=normLayer)
+        self.down1 = DownsampleBlock(64, 128, convNum=2, normLayer=normLayer)
+        self.down2 = DownsampleBlock(128, 256, convNum=2, normLayer=normLayer)
+        self.residual = nn.Sequential(*[ResidualBlock(256) for _ in range(resNum)])
+        self.up2 = UpsampleBlock(256, 128, convNum=3, normLayer=normLayer)
+        self.up1 = UpsampleBlock(128, 64, convNum=3, normLayer=normLayer)
+        self.outConv = nn.Conv2d(64, outChannel, kernel_size=3, padding=1)
+
+    def forward(self, x):
+        f1 = self.inConv(x)
+        f2 = self.down1(f1)
+        f3 = self.down2(f2)
+        r3 = self.residual(f3)
+        r2 = self.up2(r3, f2)
+        r1 = self.up1(r2, f1)
+        y = self.outConv(r1)
+        return y
+
+
+class ColorProbNet(nn.Module):
+    def __init__(self, inChannel=1, outChannel=2, with_SA=False):
+        super(ColorProbNet, self).__init__()
+        BNFunc = nn.BatchNorm2d
+        # conv1: 256
+        conv1_2 = [spectral_norm(nn.Conv2d(inChannel, 64, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv1_2 += [spectral_norm(nn.Conv2d(64, 64, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv1_2 += [BNFunc(64, affine=True)]
+        # conv2: 128
+        conv2_3 = [spectral_norm(nn.Conv2d(64, 128, 3, stride=2, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv2_3 += [spectral_norm(nn.Conv2d(128, 128, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),] 
+        conv2_3 += [spectral_norm(nn.Conv2d(128, 128, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),] 
+        conv2_3 += [BNFunc(128, affine=True)]
+        # conv3: 64
+        conv3_3 = [spectral_norm(nn.Conv2d(128, 256, 3, stride=2, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv3_3 += [spectral_norm(nn.Conv2d(256, 256, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv3_3 += [spectral_norm(nn.Conv2d(256, 256, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]  
+        conv3_3 += [BNFunc(256, affine=True)]
+        # conv4: 32
+        conv4_3 = [spectral_norm(nn.Conv2d(256, 512, 3, stride=2, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv4_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),] 
+        conv4_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv4_3 += [BNFunc(512, affine=True)]
+        # conv5: 32
+        conv5_3 = [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv5_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv5_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv5_3 += [BNFunc(512, affine=True)]
+        # conv6: 32
+        conv6_3 = [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),] 
+        conv6_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv6_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv6_3 += [BNFunc(512, affine=True),]
+        if with_SA:
+            conv6_3 += [Self_Attn(512)]
+        # conv7: 32
+        conv7_3 = [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv7_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]  
+        conv7_3 += [spectral_norm(nn.Conv2d(512, 512, 3, stride=1, padding=1)), nn.LeakyReLU(0.2, True),]
+        conv7_3 += [BNFunc(512, affine=True)]
+        # conv8: 64
+        conv8up = [nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(512, 256, 3, stride=1, padding=1),]
+        conv3short8 = [nn.Conv2d(256, 256, 3, stride=1, padding=1),]
+        conv8_3 = [nn.ReLU(True),]
+        conv8_3 += [nn.Conv2d(256, 256, 3, stride=1, padding=1), nn.ReLU(True),]
+        conv8_3 += [nn.Conv2d(256, 256, 3, stride=1, padding=1), nn.ReLU(True),]
+        conv8_3 += [BNFunc(256, affine=True),]
+        # conv9: 128
+        conv9up = [nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(256, 128, 3, stride=1, padding=1),]
+        conv9_2 = [nn.Conv2d(128, 128, 3, stride=1, padding=1), nn.ReLU(True),]
+        conv9_2 += [BNFunc(128, affine=True)]
+        # conv10: 64
+        conv10up = [nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(128, 64, 3, stride=1, padding=1),]
+        conv10_2 = [nn.ReLU(True),]
+        conv10_2 += [nn.Conv2d(64, outChannel, 3, stride=1, padding=1), nn.ReLU(True),]
+        
+        self.conv1_2 = nn.Sequential(*conv1_2)
+        self.conv2_3 = nn.Sequential(*conv2_3)
+        self.conv3_3 = nn.Sequential(*conv3_3)
+        self.conv4_3 = nn.Sequential(*conv4_3)
+        self.conv5_3 = nn.Sequential(*conv5_3)
+        self.conv6_3 = nn.Sequential(*conv6_3)
+        self.conv7_3 = nn.Sequential(*conv7_3)
+        self.conv8up = nn.Sequential(*conv8up)
+        self.conv3short8 = nn.Sequential(*conv3short8)
+        self.conv8_3 = nn.Sequential(*conv8_3)
+        self.conv9up = nn.Sequential(*conv9up)
+        self.conv9_2 = nn.Sequential(*conv9_2)
+        self.conv10up = nn.Sequential(*conv10up)
+        self.conv10_2 = nn.Sequential(*conv10_2)
+        # claffificaton output
+        #self.model_class = nn.Sequential(*[nn.Conv2d(256, 313, kernel_size=1, padding=0, stride=1),])
+
+    def forward(self, input_grays):
+        f1_2 = self.conv1_2(input_grays)
+        f2_3 = self.conv2_3(f1_2)
+        f3_3 = self.conv3_3(f2_3)
+        f4_3 = self.conv4_3(f3_3)
+        f5_3 = self.conv5_3(f4_3)
+        f6_3 = self.conv6_3(f5_3)
+        f7_3 = self.conv7_3(f6_3)
+        f8_up = self.conv8up(f7_3) + self.conv3short8(f3_3)
+        f8_3 = self.conv8_3(f8_up)
+        f9_up = self.conv9up(f8_3)
+        f9_2 = self.conv9_2(f9_up)
+        f10_up = self.conv10up(f9_2)
+        f10_2 = self.conv10_2(f10_up)
+        out_feats = f10_2
+        #out_probs = self.model_class(f8_3)
+        return out_feats
+
+
+
+def conv(batchNorm, in_planes, out_planes, kernel_size=3, stride=1):
+    if batchNorm:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=False),
+            nn.BatchNorm2d(out_planes),
+            nn.LeakyReLU(0.1)
+        )
+    else:
+        return nn.Sequential(
+            nn.Conv2d(in_planes, out_planes, kernel_size=kernel_size, stride=stride, padding=(kernel_size-1)//2, bias=True),
+            nn.LeakyReLU(0.1)
+        )
+
+
+def deconv(in_planes, out_planes):
+    return nn.Sequential(
+        nn.ConvTranspose2d(in_planes, out_planes, kernel_size=4, stride=2, padding=1, bias=True),
+        nn.LeakyReLU(0.1)
+    )
+
+class SpixelNet(nn.Module):
+    def __init__(self, inChannel=3, outChannel=9, batchNorm=True):
+        super(SpixelNet,self).__init__()
+        self.batchNorm = batchNorm
+        self.conv0a = conv(self.batchNorm, inChannel, 16, kernel_size=3)
+        self.conv0b = conv(self.batchNorm, 16, 16, kernel_size=3)
+        self.conv1a = conv(self.batchNorm, 16, 32, kernel_size=3, stride=2)
+        self.conv1b = conv(self.batchNorm, 32, 32, kernel_size=3)
+        self.conv2a = conv(self.batchNorm, 32, 64, kernel_size=3, stride=2)
+        self.conv2b = conv(self.batchNorm, 64, 64, kernel_size=3)
+        self.conv3a = conv(self.batchNorm, 64, 128, kernel_size=3, stride=2)
+        self.conv3b = conv(self.batchNorm, 128, 128, kernel_size=3)
+        self.conv4a = conv(self.batchNorm, 128, 256, kernel_size=3, stride=2)
+        self.conv4b = conv(self.batchNorm, 256, 256, kernel_size=3)
+        self.deconv3 = deconv(256, 128)
+        self.conv3_1 = conv(self.batchNorm, 256, 128)
+        self.deconv2 = deconv(128, 64)
+        self.conv2_1 = conv(self.batchNorm, 128, 64)
+        self.deconv1 = deconv(64, 32)
+        self.conv1_1 = conv(self.batchNorm, 64, 32)
+        self.deconv0 = deconv(32, 16)
+        self.conv0_1 = conv(self.batchNorm, 32, 16)
+        self.pred_mask0 = nn.Conv2d(16, outChannel, kernel_size=3, stride=1, padding=1, bias=True)
+        self.softmax = nn.Softmax(1)
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+                init.kaiming_normal_(m.weight, 0.1)
+                if m.bias is not None:
+                    init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                init.constant_(m.weight, 1)
+                init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        out1 = self.conv0b(self.conv0a(x))    #5*5
+        out2 = self.conv1b(self.conv1a(out1)) #11*11
+        out3 = self.conv2b(self.conv2a(out2)) #23*23
+        out4 = self.conv3b(self.conv3a(out3)) #47*47
+        out5 = self.conv4b(self.conv4a(out4)) #95*95
+        out_deconv3 = self.deconv3(out5)
+        concat3 = torch.cat((out4, out_deconv3), 1)
+        out_conv3_1 = self.conv3_1(concat3)
+        out_deconv2 = self.deconv2(out_conv3_1)
+        concat2 = torch.cat((out3, out_deconv2), 1)
+        out_conv2_1 = self.conv2_1(concat2)
+        out_deconv1 = self.deconv1(out_conv2_1)
+        concat1 = torch.cat((out2, out_deconv1), 1)
+        out_conv1_1 = self.conv1_1(concat1)
+        out_deconv0 = self.deconv0(out_conv1_1)
+        concat0 = torch.cat((out1, out_deconv0), 1)
+        out_conv0_1 = self.conv0_1(concat0)
+        mask0 = self.pred_mask0(out_conv0_1)
+        prob0 = self.softmax(mask0)
+        return prob0
+
+
+
+## VGG architecter, used for the perceptual loss using a pretrained VGG network
+class VGG19(torch.nn.Module):
+    def __init__(self, requires_grad=False, local_pretrained_path='checkpoints/vgg19.pth'):
+        super().__init__()
+        #vgg_pretrained_features = torchvision.models.vgg19(pretrained=True).features
+        model = torchvision.models.vgg19()
+        model.load_state_dict(torch.load(local_pretrained_path))
+        vgg_pretrained_features = model.features
+
+        self.slice1 = torch.nn.Sequential()
+        self.slice2 = torch.nn.Sequential()
+        self.slice3 = torch.nn.Sequential()
+        self.slice4 = torch.nn.Sequential()
+        self.slice5 = torch.nn.Sequential()
+        for x in range(2):
+            self.slice1.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(2, 7):
+            self.slice2.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(7, 12):
+            self.slice3.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(12, 21):
+            self.slice4.add_module(str(x), vgg_pretrained_features[x])
+        for x in range(21, 30):
+            self.slice5.add_module(str(x), vgg_pretrained_features[x])
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def forward(self, X):
+        h_relu1 = self.slice1(X)
+        h_relu2 = self.slice2(h_relu1)        
+        h_relu3 = self.slice3(h_relu2)        
+        h_relu4 = self.slice4(h_relu3)        
+        h_relu5 = self.slice5(h_relu4)                
+        out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
+        return out