Spaces:
Running
Running
| 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 |