|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
import torchvision.models as M |
|
import math |
|
from torch import Tensor |
|
from torch.nn import Parameter |
|
|
|
'''https://github.com/orashi/AlacGAN/blob/master/models/standard.py''' |
|
|
|
def l2normalize(v, eps=1e-12): |
|
return v / (v.norm() + eps) |
|
|
|
|
|
class SpectralNorm(nn.Module): |
|
def __init__(self, module, name='weight', power_iterations=1): |
|
super(SpectralNorm, self).__init__() |
|
self.module = module |
|
self.name = name |
|
self.power_iterations = power_iterations |
|
if not self._made_params(): |
|
self._make_params() |
|
|
|
def _update_u_v(self): |
|
u = getattr(self.module, self.name + "_u") |
|
v = getattr(self.module, self.name + "_v") |
|
w = getattr(self.module, self.name + "_bar") |
|
|
|
height = w.data.shape[0] |
|
for _ in range(self.power_iterations): |
|
v.data = l2normalize(torch.mv(torch.t(w.view(height,-1).data), u.data)) |
|
u.data = l2normalize(torch.mv(w.view(height,-1).data, v.data)) |
|
|
|
|
|
sigma = u.dot(w.view(height, -1).mv(v)) |
|
setattr(self.module, self.name, w / sigma.expand_as(w)) |
|
|
|
def _made_params(self): |
|
try: |
|
u = getattr(self.module, self.name + "_u") |
|
v = getattr(self.module, self.name + "_v") |
|
w = getattr(self.module, self.name + "_bar") |
|
return True |
|
except AttributeError: |
|
return False |
|
|
|
|
|
def _make_params(self): |
|
w = getattr(self.module, self.name) |
|
height = w.data.shape[0] |
|
width = w.view(height, -1).data.shape[1] |
|
|
|
u = Parameter(w.data.new(height).normal_(0, 1), requires_grad=False) |
|
v = Parameter(w.data.new(width).normal_(0, 1), requires_grad=False) |
|
u.data = l2normalize(u.data) |
|
v.data = l2normalize(v.data) |
|
w_bar = Parameter(w.data) |
|
|
|
del self.module._parameters[self.name] |
|
|
|
self.module.register_parameter(self.name + "_u", u) |
|
self.module.register_parameter(self.name + "_v", v) |
|
self.module.register_parameter(self.name + "_bar", w_bar) |
|
|
|
|
|
def forward(self, *args): |
|
self._update_u_v() |
|
return self.module.forward(*args) |
|
|
|
class Selayer(nn.Module): |
|
def __init__(self, inplanes): |
|
super(Selayer, self).__init__() |
|
self.global_avgpool = nn.AdaptiveAvgPool2d(1) |
|
self.conv1 = nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1) |
|
self.conv2 = nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1) |
|
self.relu = nn.ReLU(inplace=True) |
|
self.sigmoid = nn.Sigmoid() |
|
|
|
def forward(self, x): |
|
out = self.global_avgpool(x) |
|
out = self.conv1(out) |
|
out = self.relu(out) |
|
out = self.conv2(out) |
|
out = self.sigmoid(out) |
|
|
|
return x * out |
|
|
|
class SelayerSpectr(nn.Module): |
|
def __init__(self, inplanes): |
|
super(SelayerSpectr, self).__init__() |
|
self.global_avgpool = nn.AdaptiveAvgPool2d(1) |
|
self.conv1 = SpectralNorm(nn.Conv2d(inplanes, inplanes // 16, kernel_size=1, stride=1)) |
|
self.conv2 = SpectralNorm(nn.Conv2d(inplanes // 16, inplanes, kernel_size=1, stride=1)) |
|
self.relu = nn.ReLU(inplace=True) |
|
self.sigmoid = nn.Sigmoid() |
|
|
|
def forward(self, x): |
|
out = self.global_avgpool(x) |
|
out = self.conv1(out) |
|
out = self.relu(out) |
|
out = self.conv2(out) |
|
out = self.sigmoid(out) |
|
|
|
return x * out |
|
|
|
class ResNeXtBottleneck(nn.Module): |
|
def __init__(self, in_channels=256, out_channels=256, stride=1, cardinality=32, dilate=1): |
|
super(ResNeXtBottleneck, self).__init__() |
|
D = out_channels // 2 |
|
self.out_channels = out_channels |
|
self.conv_reduce = nn.Conv2d(in_channels, D, kernel_size=1, stride=1, padding=0, bias=False) |
|
self.conv_conv = nn.Conv2d(D, D, kernel_size=2 + stride, stride=stride, padding=dilate, dilation=dilate, |
|
groups=cardinality, |
|
bias=False) |
|
self.conv_expand = nn.Conv2d(D, out_channels, kernel_size=1, stride=1, padding=0, bias=False) |
|
self.shortcut = nn.Sequential() |
|
if stride != 1: |
|
self.shortcut.add_module('shortcut', |
|
nn.AvgPool2d(2, stride=2)) |
|
|
|
self.selayer = Selayer(out_channels) |
|
|
|
def forward(self, x): |
|
bottleneck = self.conv_reduce.forward(x) |
|
bottleneck = F.leaky_relu(bottleneck, 0.2, True) |
|
bottleneck = self.conv_conv.forward(bottleneck) |
|
bottleneck = F.leaky_relu(bottleneck, 0.2, True) |
|
bottleneck = self.conv_expand.forward(bottleneck) |
|
bottleneck = self.selayer(bottleneck) |
|
|
|
x = self.shortcut.forward(x) |
|
return x + bottleneck |
|
|
|
class SpectrResNeXtBottleneck(nn.Module): |
|
def __init__(self, in_channels=256, out_channels=256, stride=1, cardinality=32, dilate=1): |
|
super(SpectrResNeXtBottleneck, self).__init__() |
|
D = out_channels // 2 |
|
self.out_channels = out_channels |
|
self.conv_reduce = SpectralNorm(nn.Conv2d(in_channels, D, kernel_size=1, stride=1, padding=0, bias=False)) |
|
self.conv_conv = SpectralNorm(nn.Conv2d(D, D, kernel_size=2 + stride, stride=stride, padding=dilate, dilation=dilate, |
|
groups=cardinality, |
|
bias=False)) |
|
self.conv_expand = SpectralNorm(nn.Conv2d(D, out_channels, kernel_size=1, stride=1, padding=0, bias=False)) |
|
self.shortcut = nn.Sequential() |
|
if stride != 1: |
|
self.shortcut.add_module('shortcut', |
|
nn.AvgPool2d(2, stride=2)) |
|
|
|
self.selayer = SelayerSpectr(out_channels) |
|
|
|
def forward(self, x): |
|
bottleneck = self.conv_reduce.forward(x) |
|
bottleneck = F.leaky_relu(bottleneck, 0.2, True) |
|
bottleneck = self.conv_conv.forward(bottleneck) |
|
bottleneck = F.leaky_relu(bottleneck, 0.2, True) |
|
bottleneck = self.conv_expand.forward(bottleneck) |
|
bottleneck = self.selayer(bottleneck) |
|
|
|
x = self.shortcut.forward(x) |
|
return x + bottleneck |
|
|
|
class FeatureConv(nn.Module): |
|
def __init__(self, input_dim=512, output_dim=512): |
|
super(FeatureConv, self).__init__() |
|
|
|
no_bn = True |
|
|
|
seq = [] |
|
seq.append(nn.Conv2d(input_dim, output_dim, kernel_size=3, stride=1, padding=1, bias=False)) |
|
if not no_bn: seq.append(nn.BatchNorm2d(output_dim)) |
|
seq.append(nn.ReLU(inplace=True)) |
|
seq.append(nn.Conv2d(output_dim, output_dim, kernel_size=3, stride=2, padding=1, bias=False)) |
|
if not no_bn: seq.append(nn.BatchNorm2d(output_dim)) |
|
seq.append(nn.ReLU(inplace=True)) |
|
seq.append(nn.Conv2d(output_dim, output_dim, kernel_size=3, stride=1, padding=1, bias=False)) |
|
seq.append(nn.ReLU(inplace=True)) |
|
|
|
self.network = nn.Sequential(*seq) |
|
|
|
def forward(self, x): |
|
return self.network(x) |
|
|
|
class Generator(nn.Module): |
|
def __init__(self, ngf=64): |
|
super(Generator, self).__init__() |
|
|
|
self.feature_conv = FeatureConv() |
|
|
|
self.to0 = self._make_encoder_block_first(6, 32) |
|
self.to1 = self._make_encoder_block(32, 64) |
|
self.to2 = self._make_encoder_block(64, 128) |
|
self.to3 = self._make_encoder_block(128, 256) |
|
self.to4 = self._make_encoder_block(256, 512) |
|
|
|
self.deconv_for_decoder = nn.Sequential( |
|
nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1), |
|
nn.LeakyReLU(0.2), |
|
nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1), |
|
nn.LeakyReLU(0.2), |
|
nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1), |
|
nn.LeakyReLU(0.2), |
|
nn.ConvTranspose2d(32, 3, 3, stride=1, padding=1, output_padding=0), |
|
nn.Tanh(), |
|
) |
|
|
|
tunnel4 = nn.Sequential(*[ResNeXtBottleneck(ngf * 8, ngf * 8, cardinality=32, dilate=1) for _ in range(20)]) |
|
|
|
self.tunnel4 = nn.Sequential(nn.Conv2d(ngf * 8 + 512, ngf * 8, kernel_size=3, stride=1, padding=1), |
|
nn.LeakyReLU(0.2, True), |
|
tunnel4, |
|
nn.Conv2d(ngf * 8, ngf * 4 * 4, kernel_size=3, stride=1, padding=1), |
|
nn.PixelShuffle(2), |
|
nn.LeakyReLU(0.2, True) |
|
) |
|
|
|
depth = 2 |
|
tunnel = [ResNeXtBottleneck(ngf * 4, ngf * 4, cardinality=32, dilate=1) for _ in range(depth)] |
|
tunnel += [ResNeXtBottleneck(ngf * 4, ngf * 4, cardinality=32, dilate=2) for _ in range(depth)] |
|
tunnel += [ResNeXtBottleneck(ngf * 4, ngf * 4, cardinality=32, dilate=4) for _ in range(depth)] |
|
tunnel += [ResNeXtBottleneck(ngf * 4, ngf * 4, cardinality=32, dilate=2), |
|
ResNeXtBottleneck(ngf * 4, ngf * 4, cardinality=32, dilate=1)] |
|
tunnel3 = nn.Sequential(*tunnel) |
|
|
|
self.tunnel3 = nn.Sequential(nn.Conv2d(ngf * 8, ngf * 4, kernel_size=3, stride=1, padding=1), |
|
nn.LeakyReLU(0.2, True), |
|
tunnel3, |
|
nn.Conv2d(ngf * 4, ngf * 2 * 4, kernel_size=3, stride=1, padding=1), |
|
nn.PixelShuffle(2), |
|
nn.LeakyReLU(0.2, True) |
|
) |
|
|
|
tunnel = [ResNeXtBottleneck(ngf * 2, ngf * 2, cardinality=32, dilate=1) for _ in range(depth)] |
|
tunnel += [ResNeXtBottleneck(ngf * 2, ngf * 2, cardinality=32, dilate=2) for _ in range(depth)] |
|
tunnel += [ResNeXtBottleneck(ngf * 2, ngf * 2, cardinality=32, dilate=4) for _ in range(depth)] |
|
tunnel += [ResNeXtBottleneck(ngf * 2, ngf * 2, cardinality=32, dilate=2), |
|
ResNeXtBottleneck(ngf * 2, ngf * 2, cardinality=32, dilate=1)] |
|
tunnel2 = nn.Sequential(*tunnel) |
|
|
|
self.tunnel2 = nn.Sequential(nn.Conv2d(ngf * 4, ngf * 2, kernel_size=3, stride=1, padding=1), |
|
nn.LeakyReLU(0.2, True), |
|
tunnel2, |
|
nn.Conv2d(ngf * 2, ngf * 4, kernel_size=3, stride=1, padding=1), |
|
nn.PixelShuffle(2), |
|
nn.LeakyReLU(0.2, True) |
|
) |
|
|
|
tunnel = [ResNeXtBottleneck(ngf, ngf, cardinality=16, dilate=1)] |
|
tunnel += [ResNeXtBottleneck(ngf, ngf, cardinality=16, dilate=2)] |
|
tunnel += [ResNeXtBottleneck(ngf, ngf, cardinality=16, dilate=4)] |
|
tunnel += [ResNeXtBottleneck(ngf, ngf, cardinality=16, dilate=2), |
|
ResNeXtBottleneck(ngf, ngf, cardinality=16, dilate=1)] |
|
tunnel1 = nn.Sequential(*tunnel) |
|
|
|
self.tunnel1 = nn.Sequential(nn.Conv2d(ngf * 2, ngf, kernel_size=3, stride=1, padding=1), |
|
nn.LeakyReLU(0.2, True), |
|
tunnel1, |
|
nn.Conv2d(ngf, ngf * 2, kernel_size=3, stride=1, padding=1), |
|
nn.PixelShuffle(2), |
|
nn.LeakyReLU(0.2, True) |
|
) |
|
|
|
self.exit = nn.Conv2d(ngf, 3, kernel_size=3, stride=1, padding=1) |
|
|
|
|
|
def _make_encoder_block(self, inplanes, planes): |
|
return nn.Sequential( |
|
nn.Conv2d(inplanes, planes, 3, 2, 1), |
|
nn.LeakyReLU(0.2), |
|
nn.Conv2d(planes, planes, 3, 1, 1), |
|
nn.LeakyReLU(0.2), |
|
) |
|
|
|
def _make_encoder_block_first(self, inplanes, planes): |
|
return nn.Sequential( |
|
nn.Conv2d(inplanes, planes, 3, 1, 1), |
|
nn.LeakyReLU(0.2), |
|
nn.Conv2d(planes, planes, 3, 1, 1), |
|
nn.LeakyReLU(0.2), |
|
) |
|
|
|
def forward(self, sketch, sketch_feat): |
|
|
|
x0 = self.to0(sketch) |
|
x1 = self.to1(x0) |
|
x2 = self.to2(x1) |
|
x3 = self.to3(x2) |
|
x4 = self.to4(x3) |
|
|
|
sketch_feat = self.feature_conv(sketch_feat) |
|
|
|
out = self.tunnel4(torch.cat([x4, sketch_feat], 1)) |
|
|
|
|
|
|
|
|
|
x = self.tunnel3(torch.cat([out, x3], 1)) |
|
x = self.tunnel2(torch.cat([x, x2], 1)) |
|
x = self.tunnel1(torch.cat([x, x1], 1)) |
|
x = torch.tanh(self.exit(torch.cat([x, x0], 1))) |
|
|
|
decoder_output = self.deconv_for_decoder(out) |
|
|
|
return x, decoder_output |
|
''' |
|
class Colorizer(nn.Module): |
|
def __init__(self, extractor_path = 'model/model.pth'): |
|
super(Colorizer, self).__init__() |
|
|
|
self.generator = Generator() |
|
self.extractor = se_resnext_half(dump_path=extractor_path, num_classes=370, input_channels=1) |
|
|
|
def extractor_eval(self): |
|
for param in self.extractor.parameters(): |
|
param.requires_grad = False |
|
|
|
def extractor_train(self): |
|
for param in extractor.parameters(): |
|
param.requires_grad = True |
|
|
|
def forward(self, x, extractor_grad = False): |
|
|
|
if extractor_grad: |
|
features = self.extractor(x[:, 0:1]) |
|
else: |
|
with torch.no_grad(): |
|
features = self.extractor(x[:, 0:1]).detach() |
|
|
|
fake, guide = self.generator(x, features) |
|
|
|
return fake, guide |
|
''' |
|
|
|
class Colorizer(nn.Module): |
|
def __init__(self, generator_model, extractor_model): |
|
super(Colorizer, self).__init__() |
|
|
|
self.generator = generator_model |
|
self.extractor = extractor_model |
|
|
|
def load_generator_weights(self, gen_weights): |
|
self.generator.load_state_dict(gen_weights) |
|
|
|
def load_extractor_weights(self, ext_weights): |
|
self.extractor.load_state_dict(ext_weights) |
|
|
|
def extractor_eval(self): |
|
for param in self.extractor.parameters(): |
|
param.requires_grad = False |
|
self.extractor.eval() |
|
|
|
def extractor_train(self): |
|
for param in extractor.parameters(): |
|
param.requires_grad = True |
|
self.extractor.train() |
|
|
|
def forward(self, x, extractor_grad = False): |
|
|
|
if extractor_grad: |
|
features = self.extractor(x[:, 0:1]) |
|
else: |
|
with torch.no_grad(): |
|
features = self.extractor(x[:, 0:1]).detach() |
|
|
|
fake, guide = self.generator(x, features) |
|
|
|
return fake, guide |
|
|
|
class Discriminator(nn.Module): |
|
def __init__(self, ndf=64): |
|
super(Discriminator, self).__init__() |
|
|
|
self.feed = nn.Sequential(SpectralNorm(nn.Conv2d(3, 64, 3, 1, 1)), |
|
nn.LeakyReLU(0.2, True), |
|
SpectralNorm(nn.Conv2d(64, 64, 3, 2, 0)), |
|
nn.LeakyReLU(0.2, True), |
|
|
|
|
|
|
|
|
|
SpectrResNeXtBottleneck(ndf, ndf, cardinality=8, dilate=1), |
|
SpectrResNeXtBottleneck(ndf, ndf, cardinality=8, dilate=1, stride=2), |
|
SpectralNorm(nn.Conv2d(ndf, ndf * 2, kernel_size=1, stride=1, padding=0, bias=False)), |
|
nn.LeakyReLU(0.2, True), |
|
|
|
SpectrResNeXtBottleneck(ndf * 2, ndf * 2, cardinality=8, dilate=1), |
|
SpectrResNeXtBottleneck(ndf * 2, ndf * 2, cardinality=8, dilate=1, stride=2), |
|
SpectralNorm(nn.Conv2d(ndf * 2, ndf * 4, kernel_size=1, stride=1, padding=0, bias=False)), |
|
nn.LeakyReLU(0.2, True), |
|
|
|
SpectrResNeXtBottleneck(ndf * 4, ndf * 4, cardinality=8, dilate=1), |
|
SpectrResNeXtBottleneck(ndf * 4, ndf * 4, cardinality=8, dilate=1, stride=2), |
|
SpectralNorm(nn.Conv2d(ndf * 4, ndf * 8, kernel_size=1, stride=1, padding=1, bias=False)), |
|
nn.LeakyReLU(0.2, True), |
|
SpectrResNeXtBottleneck(ndf * 8, ndf * 8, cardinality=8, dilate=1), |
|
SpectrResNeXtBottleneck(ndf * 8, ndf * 8, cardinality=8, dilate=1, stride=2), |
|
SpectrResNeXtBottleneck(ndf * 8, ndf * 8, cardinality=8, dilate=1), |
|
SpectrResNeXtBottleneck(ndf * 8, ndf * 8, cardinality=8, dilate=1), |
|
nn.AdaptiveAvgPool2d((1, 1)) |
|
) |
|
|
|
self.out = nn.Linear(512, 1) |
|
|
|
def forward(self, color): |
|
x = self.feed(color) |
|
|
|
out = self.out(x.view(color.size(0), -1)) |
|
return out |
|
|
|
class Content(nn.Module): |
|
def __init__(self, path): |
|
super(Content, self).__init__() |
|
vgg16 = M.vgg16() |
|
vgg16.load_state_dict(torch.load(path)) |
|
vgg16.features = nn.Sequential( |
|
*list(vgg16.features.children())[:9] |
|
) |
|
self.model = vgg16.features |
|
self.register_buffer('mean', torch.FloatTensor([0.485 - 0.5, 0.456 - 0.5, 0.406 - 0.5]).view(1, 3, 1, 1)) |
|
self.register_buffer('std', torch.FloatTensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)) |
|
|
|
def forward(self, images): |
|
return self.model((images.mul(0.5) - self.mean) / self.std) |
|
|