初始化：换发型/换发色/训练发型服务
包含： - hair_service_sd: 主服务（换发型/换发色/生发，端口8801） - photo_service: LoRA调度+训练（端口32678） - hair_grow_service: 调试测试页（端口8888，含4个测试页） - 批量训练脚本（batch_train_hairstyles.py） - 发际线mask自动识别（hairline_mask.py，4种方案） - 手绘mask换发型（hair_swap_manual.py） - 文档：README.md + LARGE_FILES.md + docs/ 大文件（模型权重200G、训练数据123G）已排除，见 LARGE_FILES.md OSS/COS密钥已脱敏为环境变量，原文件备份在本地
This commit is contained in:
xsl
@@ -0,0 +1 @@
|
||||
from .generators import *
|
||||
@@ -0,0 +1,28 @@
|
||||
from .resnet_dec import ResNet_D_Dec, BasicBlock
|
||||
from .res_shortcut_dec import ResShortCut_D_Dec
|
||||
from .res_gca_dec import ResGuidedCxtAtten_Dec
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_decoder_22', 'res_gca_decoder_22']
|
||||
|
||||
|
||||
def _res_shortcut_D_dec(block, layers, **kwargs):
|
||||
model = ResShortCut_D_Dec(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D_dec(block, layers, num_class, **kwargs):
|
||||
model = ResGuidedCxtAtten_Dec(block, layers, num_class, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def res_shortcut_decoder_22(**kwargs):
|
||||
"""Constructs a resnet_encoder_14 model.
|
||||
"""
|
||||
return _res_shortcut_D_dec(BasicBlock, [2, 3, 3, 2], **kwargs)
|
||||
|
||||
|
||||
def res_gca_decoder_22(num_class=1, **kwargs):
|
||||
"""Constructs a resnet_encoder_14 model.
|
||||
"""
|
||||
return _res_gca_D_dec(BasicBlock, [2, 3, 3, 2], num_class, **kwargs)
|
||||
@@ -0,0 +1,28 @@
|
||||
from matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||
from matting.networks.decoders.res_shortcut_dec import ResShortCut_D_Dec
|
||||
|
||||
|
||||
class ResGuidedCxtAtten_Dec(ResShortCut_D_Dec):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False):
|
||||
super(ResGuidedCxtAtten_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel)
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
self.num_class = num_class
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||
im = mid_fea['image_fea']
|
||||
x = self.layer1(x) + fea5 # N x 256 x 32 x 32
|
||||
x = self.layer2(x) + fea4 # N x 128 x 64 x 64
|
||||
x, offset = self.gca(im, x, mid_fea['unknown']) # contextual attention
|
||||
x = self.layer3(x) + fea3 # N x 64 x 128 x 128
|
||||
x = self.layer4(x) + fea2 # N x 32 x 256 x 256
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x) + fea1
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, {'offset_1': mid_fea['offset_1'], 'offset_2': offset}
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
from matting.networks.decoders.resnet_dec import ResNet_D_Dec
|
||||
|
||||
|
||||
class ResShortCut_D_Dec(ResNet_D_Dec):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||
super(ResShortCut_D_Dec, self).__init__(block, layers, num_class, norm_layer, large_kernel,
|
||||
late_downsample=late_downsample)
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
fea1, fea2, fea3, fea4, fea5 = mid_fea['shortcut']
|
||||
x = self.layer1(x) + fea5
|
||||
x = self.layer2(x) + fea4
|
||||
x = self.layer3(x) + fea3
|
||||
x = self.layer4(x) + fea2
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x) + fea1
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, None
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv5x5(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""5x5 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=5, stride=stride,
|
||||
padding=2, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, upsample=None, norm_layer=None, large_kernel=False):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self.stride = stride
|
||||
conv = conv5x5 if large_kernel else conv3x3
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
if self.stride > 1:
|
||||
self.conv1 = SpectralNorm(nn.ConvTranspose2d(inplanes, inplanes, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
else:
|
||||
self.conv1 = SpectralNorm(conv(inplanes, inplanes))
|
||||
self.bn1 = norm_layer(inplanes)
|
||||
self.activation = nn.LeakyReLU(0.2, inplace=True)
|
||||
self.conv2 = SpectralNorm(conv(inplanes, planes))
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.upsample = upsample
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.upsample is not None:
|
||||
identity = self.upsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D_Dec(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, num_class=1, norm_layer=None, large_kernel=False, late_downsample=False):
|
||||
super(ResNet_D_Dec, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
self.large_kernel = large_kernel
|
||||
self.kernel_size = 5 if self.large_kernel else 3
|
||||
|
||||
self.inplanes = 512 if layers[0] > 0 else 256
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
|
||||
self.conv1 = SpectralNorm(nn.ConvTranspose2d(self.midplanes, 32, kernel_size=4, stride=2, padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.leaky_relu = nn.LeakyReLU(0.2, inplace=True)
|
||||
self.conv2 = nn.Conv2d(32, num_class, kernel_size=self.kernel_size, stride=1, padding=self.kernel_size//2)
|
||||
self.upsample = nn.UpsamplingNearest2d(scale_factor=2)
|
||||
self.tanh = nn.Tanh()
|
||||
self.layer1 = self._make_layer(block, 256, layers[0], stride=2)
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 64, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, self.midplanes, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
else:
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
upsample = None
|
||||
if stride != 1:
|
||||
upsample = nn.Sequential(
|
||||
nn.UpsamplingNearest2d(scale_factor=2),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
upsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, upsample, norm_layer, self.large_kernel)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer, large_kernel=self.large_kernel))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x, mid_fea):
|
||||
x = self.layer1(x) # N x 256 x 32 x 32
|
||||
x = self.layer2(x) # N x 128 x 64 x 64
|
||||
x = self.layer3(x) # N x 64 x 128 x 128
|
||||
x = self.layer4(x) # N x 32 x 256 x 256
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.leaky_relu(x)
|
||||
x = self.conv2(x)
|
||||
|
||||
alpha = (self.tanh(x) + 1.0) / 2.0
|
||||
|
||||
return alpha, None
|
||||
@@ -0,0 +1,39 @@
|
||||
import logging
|
||||
from .resnet_enc import ResNet_D, BasicBlock
|
||||
from .res_shortcut_enc import ResShortCut_D
|
||||
from .res_gca_enc import ResGuidedCxtAtten
|
||||
|
||||
|
||||
__all__ = ['res_shortcut_encoder_29', 'resnet_gca_encoder_29']
|
||||
|
||||
|
||||
def _res_shortcut_D(block, layers, **kwargs):
|
||||
model = ResShortCut_D(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def _res_gca_D(block, layers, **kwargs):
|
||||
model = ResGuidedCxtAtten(block, layers, **kwargs)
|
||||
return model
|
||||
|
||||
|
||||
def resnet_gca_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_29 model.
|
||||
"""
|
||||
return _res_gca_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
def res_shortcut_encoder_29(**kwargs):
|
||||
"""Constructs a resnet_encoder_25 model.
|
||||
"""
|
||||
return _res_shortcut_D(BasicBlock, [3, 4, 4, 2], **kwargs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import torch
|
||||
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s: %(message)s',
|
||||
datefmt='%m-%d %H:%M:%S')
|
||||
resnet_encoder = res_shortcut_encoder_29()
|
||||
x = torch.randn(4,6,512,512)
|
||||
z = resnet_encoder(x)
|
||||
print(z[0].shape)
|
||||
@@ -0,0 +1,97 @@
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
# from utils import CONFIG
|
||||
from matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from matting.networks.ops import GuidedCxtAtten, SpectralNorm
|
||||
|
||||
class ResGuidedCxtAtten(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResGuidedCxtAtten, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
self.guidance_head = nn.Sequential(
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(3, 16, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(16),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(16, 32, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(32),
|
||||
nn.ReflectionPad2d(1),
|
||||
SpectralNorm(nn.Conv2d(32, 128, kernel_size=3, padding=0, stride=2, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(128)
|
||||
)
|
||||
|
||||
self.gca = GuidedCxtAtten(128, 128)
|
||||
|
||||
# initialize guidance head
|
||||
for layers in range(len(self.guidance_head)):
|
||||
m = self.guidance_head[layers]
|
||||
if isinstance(m, nn.Conv2d):
|
||||
if hasattr(m, "weight_bar"):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
im_fea = self.guidance_head(x[:, :3, ...]) # downsample origin image and extract features
|
||||
# if CONFIG.model.trimap_channel == 3:
|
||||
unknown = F.interpolate(x[:, 4:5, ...], scale_factor=1/8, mode='nearest')
|
||||
# else:
|
||||
# unknown = F.interpolate(x[:,3:,...].eq(1.).float(), scale_factor=1/8, mode='nearest')
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x3, offset = self.gca(im_fea, x3, unknown) # contextual attention
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut': (fea1, fea2, fea3, fea4, fea5),
|
||||
'image_fea': im_fea,
|
||||
'unknown': unknown,
|
||||
'offset_1': offset}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from matting.networks.encoders.resnet_enc import BasicBlock
|
||||
m = ResGuidedCxtAtten(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m)
|
||||
@@ -0,0 +1,51 @@
|
||||
import torch.nn as nn
|
||||
# from utils import CONFIG
|
||||
from matting.networks.encoders.resnet_enc import ResNet_D
|
||||
from matting.networks.ops import SpectralNorm
|
||||
|
||||
class ResShortCut_D(ResNet_D):
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResShortCut_D, self).__init__(block, layers, norm_layer, late_downsample=late_downsample)
|
||||
first_inplane = 3 + 3
|
||||
self.shortcut_inplane = [first_inplane, self.midplanes, 64, 128, 256]
|
||||
self.shortcut_plane = [32, self.midplanes, 64, 128, 256]
|
||||
|
||||
self.shortcut = nn.ModuleList()
|
||||
for stage, inplane in enumerate(self.shortcut_inplane):
|
||||
self.shortcut.append(self._make_shortcut(inplane, self.shortcut_plane[stage]))
|
||||
|
||||
def _make_shortcut(self, inplane, planes):
|
||||
return nn.Sequential(
|
||||
SpectralNorm(nn.Conv2d(inplane, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes),
|
||||
SpectralNorm(nn.Conv2d(planes, planes, kernel_size=3, padding=1, bias=False)),
|
||||
nn.ReLU(inplace=True),
|
||||
self._norm_layer(planes)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
x1 = self.activation(out) # N x 32 x 256 x 256
|
||||
out = self.conv3(x1)
|
||||
out = self.bn3(out)
|
||||
out = self.activation(out)
|
||||
|
||||
x2 = self.layer1(out) # N x 64 x 128 x 128
|
||||
x3= self.layer2(x2) # N x 128 x 64 x 64
|
||||
x4 = self.layer3(x3) # N x 256 x 32 x 32
|
||||
out = self.layer_bottleneck(x4) # N x 512 x 16 x 16
|
||||
|
||||
fea1 = self.shortcut[0](x) # input image and trimap
|
||||
fea2 = self.shortcut[1](x1)
|
||||
fea3 = self.shortcut[2](x2)
|
||||
fea4 = self.shortcut[3](x3)
|
||||
fea5 = self.shortcut[4](x4)
|
||||
|
||||
return out, {'shortcut':(fea1, fea2, fea3, fea4, fea5), 'image':x[:,:3,...]}
|
||||
@@ -0,0 +1,150 @@
|
||||
import logging
|
||||
import torch.nn as nn
|
||||
from matting.networks.ops import SpectralNorm
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=False, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, norm_layer=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = SpectralNorm(conv3x3(inplanes, planes, stride))
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.conv2 = SpectralNorm(conv3x3(planes, planes))
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.activation(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
|
||||
out += identity
|
||||
out = self.activation(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class ResNet_D(nn.Module):
|
||||
"""
|
||||
Implement and pre-train on ImageNet with the tricks from
|
||||
https://arxiv.org/abs/1812.01187
|
||||
without the mix-up part.
|
||||
"""
|
||||
|
||||
def __init__(self, block, layers, norm_layer=None, late_downsample=False):
|
||||
super(ResNet_D, self).__init__()
|
||||
self.logger = logging.getLogger("Logger")
|
||||
if norm_layer is None:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
self._norm_layer = norm_layer
|
||||
|
||||
self.inplanes = 64
|
||||
self.late_downsample = late_downsample
|
||||
self.midplanes = 64 if late_downsample else 32
|
||||
self.start_stride = [1, 2, 1, 2] if late_downsample else [2, 1, 2, 1]
|
||||
self.conv1 = SpectralNorm(nn.Conv2d(3 + 3, 32, kernel_size=3,
|
||||
stride=self.start_stride[0], padding=1, bias=False))
|
||||
self.conv2 = SpectralNorm(nn.Conv2d(32, self.midplanes, kernel_size=3, stride=self.start_stride[1], padding=1,
|
||||
bias=False))
|
||||
self.conv3 = SpectralNorm(nn.Conv2d(self.midplanes, self.inplanes, kernel_size=3, stride=self.start_stride[2],
|
||||
padding=1, bias=False))
|
||||
self.bn1 = norm_layer(32)
|
||||
self.bn2 = norm_layer(self.midplanes)
|
||||
self.bn3 = norm_layer(self.inplanes)
|
||||
self.activation = nn.ReLU(inplace=True)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0], stride=self.start_stride[3])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer_bottleneck = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
nn.init.xavier_uniform_(m.weight_bar)
|
||||
elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
# Zero-initialize the last BN in each residual branch,
|
||||
# so that the residual branch starts with zeros, and each residual block behaves like an identity.
|
||||
# This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
|
||||
|
||||
for m in self.modules():
|
||||
if isinstance(m, BasicBlock):
|
||||
nn.init.constant_(m.bn2.weight, 0)
|
||||
|
||||
self.logger.debug("encoder conv1 weight shape: {}".format(str(self.conv1.module.weight_bar.data.shape)))
|
||||
self.conv1.module.weight_bar.data[:,3:,:,:] = 0
|
||||
|
||||
self.logger.debug(self)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
if blocks == 0:
|
||||
return nn.Sequential(nn.Identity())
|
||||
norm_layer = self._norm_layer
|
||||
downsample = None
|
||||
if stride != 1:
|
||||
downsample = nn.Sequential(
|
||||
nn.AvgPool2d(2, stride),
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
elif self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
SpectralNorm(conv1x1(self.inplanes, planes * block.expansion, stride)),
|
||||
norm_layer(planes * block.expansion),
|
||||
)
|
||||
|
||||
layers = [block(self.inplanes, planes, stride, downsample, norm_layer)]
|
||||
self.inplanes = planes * block.expansion
|
||||
for _ in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes, norm_layer=norm_layer))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.activation(x)
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x1 = self.activation(x) # N x 32 x 256 x 256
|
||||
x = self.conv3(x1)
|
||||
x = self.bn3(x)
|
||||
x2 = self.activation(x) # N x 64 x 128 x 128
|
||||
|
||||
x3 = self.layer1(x2) # N x 64 x 128 x 128
|
||||
x4 = self.layer2(x3) # N x 128 x 64 x 64
|
||||
x5 = self.layer3(x4) # N x 256 x 32 x 32
|
||||
x = self.layer_bottleneck(x5) # N x 512 x 16 x 16
|
||||
|
||||
return x, (x1, x2, x3, x4, x5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
m = ResNet_D(BasicBlock, [3, 4, 4, 2])
|
||||
for m in m.modules():
|
||||
print(m._get_name())
|
||||
@@ -0,0 +1,60 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
# from utils import CONFIG
|
||||
from matting.networks import encoders, decoders
|
||||
|
||||
|
||||
class Generator(nn.Module):
|
||||
def __init__(self, encoder, decoder, num_class=1):
|
||||
|
||||
super(Generator, self).__init__()
|
||||
|
||||
if encoder not in encoders.__all__:
|
||||
raise NotImplementedError("Unknown Encoder {}".format(encoder))
|
||||
self.encoder = encoders.__dict__[encoder]()
|
||||
|
||||
if decoder not in decoders.__all__:
|
||||
raise NotImplementedError("Unknown Decoder {}".format(decoder))
|
||||
self.decoder = decoders.__dict__[decoder](num_class)
|
||||
|
||||
def forward(self, image, trimap):
|
||||
inp = torch.cat((image, trimap), dim=1)
|
||||
embedding, mid_fea = self.encoder(inp)
|
||||
alpha, info_dict = self.decoder(embedding, mid_fea)
|
||||
|
||||
return alpha, info_dict
|
||||
|
||||
|
||||
def get_generator(encoder, decoder, num_class=1):
|
||||
generator = Generator(encoder=encoder, decoder=decoder, num_class=num_class)
|
||||
return generator
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
import time
|
||||
# generator = get_generator(encoder=CONFIG.model.arch.encoder, decoder=CONFIG.model.arch.decoder).cuda().train()
|
||||
batch_size = 12
|
||||
# generator.eval()
|
||||
n_eval = 10
|
||||
# pre run the model
|
||||
# with torch.no_grad():
|
||||
# for i in range(2):
|
||||
# x = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||
# y = torch.rand(batch_size, 3, 512, 512, device=device)
|
||||
# z = generator(x,y)
|
||||
# test without GPU IO
|
||||
|
||||
# x = torch.zeros(batch_size, 3, 512, 512, device=device)
|
||||
# y = torch.zeros(batch_size, 1, 512, 512, device=device)
|
||||
x = torch.randn(batch_size, 3, 512, 512)
|
||||
y = torch.randn(batch_size, 3, 512, 512)
|
||||
t = time.time()
|
||||
# with torch.no_grad():
|
||||
# for i in range(n_eval):
|
||||
# a = generator(x.cuda(),y.cuda())
|
||||
# torch.cuda.synchronize()
|
||||
# print(generator.__class__.__name__, 'With IO \t', f'{(time.time() - t)/n_eval/batch_size:.5f} s')
|
||||
# print(generator.__class__.__name__, 'FPS \t\t', f'{1/((time.time() - t)/n_eval/batch_size):.5f} s')
|
||||
# for n, p in generator.named_parameters():
|
||||
# print(n)
|
||||
@@ -0,0 +1,256 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import Parameter
|
||||
from torch.autograd import Variable
|
||||
from torch.nn import functional as F
|
||||
|
||||
|
||||
def l2normalize(v, eps=1e-12):
|
||||
return v / (v.norm() + eps)
|
||||
|
||||
|
||||
class SpectralNorm(nn.Module):
|
||||
"""
|
||||
Based on https://github.com/heykeetae/Self-Attention-GAN/blob/master/spectral.py
|
||||
and add _noupdate_u_v() for evaluation
|
||||
"""
|
||||
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 _noupdate_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]
|
||||
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):
|
||||
# if torch.is_grad_enabled() and self.module.training:
|
||||
if self.module.training:
|
||||
self._update_u_v()
|
||||
else:
|
||||
self._noupdate_u_v()
|
||||
return self.module.forward(*args)
|
||||
|
||||
|
||||
class GuidedCxtAtten(nn.Module):
|
||||
# based on https://github.com/nbei/Deep-Flow-Guided-Video-Inpainting/blob/a6fe298fec502bfd9cbc64eb01e39f78a3262a59/models/DeepFill_Models/ops.py#L210
|
||||
def __init__(self, out_channels, guidance_channels, rate=2):
|
||||
super(GuidedCxtAtten, self).__init__()
|
||||
self.rate = rate
|
||||
self.padding = nn.ReflectionPad2d(1)
|
||||
self.up_sample = nn.Upsample(scale_factor=self.rate, mode='nearest')
|
||||
|
||||
self.guidance_conv = nn.Conv2d(in_channels=guidance_channels, out_channels=guidance_channels//2,
|
||||
kernel_size=1, stride=1, padding=0)
|
||||
|
||||
self.W = nn.Sequential(
|
||||
nn.Conv2d(in_channels=out_channels, out_channels=out_channels,
|
||||
kernel_size=1, stride=1, padding=0, bias=False),
|
||||
nn.BatchNorm2d(out_channels)
|
||||
)
|
||||
|
||||
nn.init.xavier_uniform_(self.guidance_conv.weight)
|
||||
nn.init.constant_(self.guidance_conv.bias, 0)
|
||||
nn.init.xavier_uniform_(self.W[0].weight)
|
||||
nn.init.constant_(self.W[1].weight, 1e-3)
|
||||
nn.init.constant_(self.W[1].bias, 0)
|
||||
|
||||
def forward(self, f, alpha, unknown=None, ksize=3, stride=1, fuse_k=3, softmax_scale=1., training=True):
|
||||
|
||||
f = self.guidance_conv(f)
|
||||
# get shapes
|
||||
raw_int_fs = list(f.size()) # N x 64 x 64 x 64
|
||||
raw_int_alpha = list(alpha.size()) # N x 128 x 64 x 64
|
||||
|
||||
# extract patches from background with stride and rate
|
||||
kernel = 2*self.rate
|
||||
alpha_w = self.extract_patches(alpha, kernel=kernel, stride=self.rate)
|
||||
alpha_w = alpha_w.permute(0, 2, 3, 4, 5, 1)
|
||||
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], raw_int_alpha[2] // self.rate, raw_int_alpha[3] // self.rate, -1)
|
||||
alpha_w = alpha_w.contiguous().view(raw_int_alpha[0], -1, kernel, kernel, raw_int_alpha[1])
|
||||
alpha_w = alpha_w.permute(0, 1, 4, 2, 3)
|
||||
|
||||
f = F.interpolate(f, scale_factor=1/self.rate, mode='nearest')
|
||||
|
||||
fs = f.size() # B x 64 x 32 x 32
|
||||
f_groups = torch.split(f, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
|
||||
# from b(B*H*W*C) to w(b*k*k*c*h*w)
|
||||
int_fs = list(fs)
|
||||
w = self.extract_patches(f)
|
||||
w = w.permute(0, 2, 3, 4, 5, 1)
|
||||
w = w.contiguous().view(raw_int_fs[0], raw_int_fs[2] // self.rate, raw_int_fs[3] // self.rate, -1)
|
||||
w = w.contiguous().view(raw_int_fs[0], -1, ksize, ksize, raw_int_fs[1])
|
||||
w = w.permute(0, 1, 4, 2, 3)
|
||||
# process mask
|
||||
|
||||
if unknown is not None:
|
||||
unknown = unknown.clone()
|
||||
unknown = F.interpolate(unknown, scale_factor=1/self.rate, mode='nearest')
|
||||
assert unknown.size(2) == f.size(2), "mask should have same size as f at dim 2,3"
|
||||
unknown_mean = unknown.mean(dim=[2,3])
|
||||
known_mean = 1 - unknown_mean
|
||||
unknown_scale = torch.clamp(torch.sqrt(unknown_mean / known_mean), 0.1, 10).to(alpha)
|
||||
known_scale = torch.clamp(torch.sqrt(known_mean / unknown_mean), 0.1, 10).to(alpha)
|
||||
softmax_scale = torch.cat([unknown_scale, known_scale], dim=1)
|
||||
else:
|
||||
unknown = torch.ones([fs[0], 1, fs[2], fs[3]]).to(alpha)
|
||||
softmax_scale = torch.FloatTensor([softmax_scale, softmax_scale]).view(1,2).repeat(fs[0],1).to(alpha)
|
||||
|
||||
m = self.extract_patches(unknown)
|
||||
|
||||
m = m.permute(0, 2, 3, 4, 5, 1)
|
||||
m = m.contiguous().view(raw_int_fs[0], raw_int_fs[2]//self.rate, raw_int_fs[3]//self.rate, -1)
|
||||
m = m.contiguous().view(raw_int_fs[0], -1, ksize, ksize)
|
||||
|
||||
m = self.reduce_mean(m) # smoothing, maybe
|
||||
# mask out the
|
||||
mm = m.gt(0.).float() # (N, 32*32, 1, 1)
|
||||
|
||||
# the correlation with itself should be 0
|
||||
self_mask = F.one_hot(torch.arange(fs[2] * fs[3]).view(fs[2], fs[3]).contiguous().to(alpha).long(),
|
||||
num_classes=int_fs[2] * int_fs[3])
|
||||
self_mask = self_mask.permute(2, 0, 1).view(1, fs[2] * fs[3], fs[2], fs[3]).float() * (-1e4)
|
||||
|
||||
w_groups = torch.split(w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
alpha_w_groups = torch.split(alpha_w, 1, dim=0) # Split tensors by batch dimension; tuple is returned
|
||||
mm_groups = torch.split(mm, 1, dim=0)
|
||||
scale_group = torch.split(softmax_scale, 1, dim=0)
|
||||
y = []
|
||||
offsets = []
|
||||
k = fuse_k
|
||||
y_test = []
|
||||
for xi, wi, alpha_wi, mmi, scale in zip(f_groups, w_groups, alpha_w_groups, mm_groups, scale_group):
|
||||
# conv for compare
|
||||
wi = wi[0]
|
||||
escape_NaN = Variable(torch.FloatTensor([1e-4])).to(alpha)
|
||||
wi_normed = wi / torch.max(self.l2_norm(wi), escape_NaN)
|
||||
xi = F.pad(xi, (1,1,1,1), mode='reflect')
|
||||
yi = F.conv2d(xi, wi_normed, stride=1, padding=0) # yi => (B=1, C=32*32, H=32, W=32)
|
||||
y_test.append(yi)
|
||||
# conv implementation for fuse scores to encourage large patches
|
||||
yi = yi.permute(0, 2, 3, 1)
|
||||
yi = yi.contiguous().view(1, fs[2], fs[3], fs[2] * fs[3])
|
||||
yi = yi.permute(0, 3, 1, 2) # (B=1, C=32*32, H=32, W=32)
|
||||
|
||||
# softmax to match
|
||||
# scale the correlation with predicted scale factor for known and unknown area
|
||||
yi = yi * (scale[0,0] * mmi.gt(0.).float() + scale[0,1] * mmi.le(0.).float()) # mmi => (1, 32*32, 1, 1)
|
||||
# mask itself, self-mask only applied to unknown area
|
||||
yi = yi + self_mask * mmi # self_mask: (1, 32*32, 32, 32)
|
||||
# for small input inference
|
||||
yi = F.softmax(yi, dim=1)
|
||||
|
||||
_, offset = torch.max(yi, dim=1) # argmax; index
|
||||
offset = torch.stack([offset // fs[3], offset % fs[3]], dim=1)
|
||||
|
||||
wi_center = alpha_wi[0]
|
||||
|
||||
if self.rate == 1:
|
||||
left = (kernel) // 2
|
||||
right = (kernel - 1) // 2
|
||||
yi = F.pad(yi, (left, right, left, right), mode='reflect')
|
||||
wi_center = wi_center.permute(1, 0, 2, 3)
|
||||
yi = F.conv2d(yi, wi_center, padding=0) / 4. # (B=1, C=128, H=64, W=64)
|
||||
else:
|
||||
yi = F.conv_transpose2d(yi, wi_center, stride=self.rate, padding=1) / 4. # (B=1, C=128, H=64, W=64)
|
||||
y.append(yi)
|
||||
offsets.append(offset)
|
||||
|
||||
y = torch.cat(y, dim=0) # back to the mini-batch
|
||||
y.contiguous().view(raw_int_alpha)
|
||||
offsets = torch.cat(offsets, dim=0)
|
||||
offsets = offsets.view([int_fs[0]] + [2] + int_fs[2:])
|
||||
|
||||
# # case1: visualize optical flow: minus current position
|
||||
# h_add = Variable(torch.arange(0,float(fs[2]))).to(alpha).view([1, 1, fs[2], 1])
|
||||
# h_add = h_add.expand(fs[0], 1, fs[2], fs[3])
|
||||
# w_add = Variable(torch.arange(0,float(fs[3]))).to(alpha).view([1, 1, 1, fs[3]])
|
||||
# w_add = w_add.expand(fs[0], 1, fs[2], fs[3])
|
||||
#
|
||||
# offsets = offsets - torch.cat([h_add, w_add], dim=1).long()
|
||||
|
||||
# case2: visualize absolute position
|
||||
offsets = offsets - torch.Tensor([fs[2]//2, fs[3]//2]).view(1,2,1,1).to(alpha).long()
|
||||
|
||||
y = self.W(y) + alpha
|
||||
|
||||
return y, (offsets, softmax_scale)
|
||||
|
||||
@staticmethod
|
||||
def extract_patches(x, kernel=3, stride=1):
|
||||
left =(kernel - stride + 1) // 2
|
||||
right =(kernel - stride) // 2
|
||||
x = F.pad(x, (left, right, left, right), mode='reflect')
|
||||
all_patches = x.unfold(2, kernel, stride).unfold(3, kernel, stride)
|
||||
|
||||
return all_patches
|
||||
|
||||
@staticmethod
|
||||
def reduce_mean(x):
|
||||
for i in range(4):
|
||||
if i <= 1:
|
||||
continue
|
||||
x = torch.mean(x, dim=i, keepdim=True)
|
||||
return x
|
||||
|
||||
@staticmethod
|
||||
def l2_norm(x):
|
||||
def reduce_sum(x):
|
||||
for i in range(4):
|
||||
if i == 0:
|
||||
continue
|
||||
x = torch.sum(x, dim=i, keepdim=True)
|
||||
return x
|
||||
|
||||
x = x**2
|
||||
x = reduce_sum(x)
|
||||
return torch.sqrt(x)
|
||||
Reference in New Issue
Block a user