You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
246 lines
7.7 KiB
246 lines
7.7 KiB
from torch import nn
|
|
|
|
import torch.nn.functional as F
|
|
import torch
|
|
|
|
from sync_batchnorm import SynchronizedBatchNorm2d as BatchNorm2d
|
|
|
|
|
|
def kp2gaussian(kp, spatial_size, kp_variance):
|
|
"""
|
|
Transform a keypoint into gaussian like representation
|
|
"""
|
|
mean = kp['value']
|
|
|
|
coordinate_grid = make_coordinate_grid(spatial_size, mean.type())
|
|
number_of_leading_dimensions = len(mean.shape) - 1
|
|
shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape
|
|
coordinate_grid = coordinate_grid.view(*shape)
|
|
repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1)
|
|
coordinate_grid = coordinate_grid.repeat(*repeats)
|
|
|
|
# Preprocess kp shape
|
|
shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 2)
|
|
mean = mean.view(*shape)
|
|
|
|
mean_sub = (coordinate_grid - mean)
|
|
|
|
out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance)
|
|
|
|
return out
|
|
|
|
|
|
def make_coordinate_grid(spatial_size, type):
|
|
"""
|
|
Create a meshgrid [-1,1] x [-1,1] of given spatial_size.
|
|
"""
|
|
h, w = spatial_size
|
|
x = torch.arange(w).type(type)
|
|
y = torch.arange(h).type(type)
|
|
|
|
x = (2 * (x / (w - 1)) - 1)
|
|
y = (2 * (y / (h - 1)) - 1)
|
|
|
|
yy = y.view(-1, 1).repeat(1, w)
|
|
xx = x.view(1, -1).repeat(h, 1)
|
|
|
|
meshed = torch.cat([xx.unsqueeze_(2), yy.unsqueeze_(2)], 2)
|
|
|
|
return meshed
|
|
|
|
|
|
class ResBlock2d(nn.Module):
|
|
"""
|
|
Res block, preserve spatial resolution.
|
|
"""
|
|
|
|
def __init__(self, in_features, kernel_size, padding):
|
|
super(ResBlock2d, self).__init__()
|
|
self.conv1 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.conv2 = nn.Conv2d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size,
|
|
padding=padding)
|
|
self.norm1 = BatchNorm2d(in_features, affine=True)
|
|
self.norm2 = BatchNorm2d(in_features, affine=True)
|
|
|
|
def forward(self, x):
|
|
out = self.norm1(x)
|
|
out = F.relu(out)
|
|
out = self.conv1(out)
|
|
out = self.norm2(out)
|
|
out = F.relu(out)
|
|
out = self.conv2(out)
|
|
out += x
|
|
return out
|
|
|
|
|
|
class UpBlock2d(nn.Module):
|
|
"""
|
|
Upsampling block for use in decoder.
|
|
"""
|
|
|
|
def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
|
|
super(UpBlock2d, self).__init__()
|
|
|
|
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
|
padding=padding, groups=groups)
|
|
self.norm = BatchNorm2d(out_features, affine=True)
|
|
|
|
def forward(self, x):
|
|
out = F.interpolate(x, scale_factor=2)
|
|
out = self.conv(out)
|
|
out = self.norm(out)
|
|
out = F.relu(out)
|
|
return out
|
|
|
|
|
|
class DownBlock2d(nn.Module):
|
|
"""
|
|
Downsampling block for use in encoder.
|
|
"""
|
|
|
|
def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
|
|
super(DownBlock2d, self).__init__()
|
|
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
|
|
padding=padding, groups=groups)
|
|
self.norm = BatchNorm2d(out_features, affine=True)
|
|
self.pool = nn.AvgPool2d(kernel_size=(2, 2))
|
|
|
|
def forward(self, x):
|
|
out = self.conv(x)
|
|
out = self.norm(out)
|
|
out = F.relu(out)
|
|
out = self.pool(out)
|
|
return out
|
|
|
|
|
|
class SameBlock2d(nn.Module):
|
|
"""
|
|
Simple block, preserve spatial resolution.
|
|
"""
|
|
|
|
def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1):
|
|
super(SameBlock2d, self).__init__()
|
|
self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features,
|
|
kernel_size=kernel_size, padding=padding, groups=groups)
|
|
self.norm = BatchNorm2d(out_features, affine=True)
|
|
|
|
def forward(self, x):
|
|
out = self.conv(x)
|
|
out = self.norm(out)
|
|
out = F.relu(out)
|
|
return out
|
|
|
|
|
|
class Encoder(nn.Module):
|
|
"""
|
|
Hourglass Encoder
|
|
"""
|
|
|
|
def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
|
|
super(Encoder, self).__init__()
|
|
|
|
down_blocks = []
|
|
for i in range(num_blocks):
|
|
down_blocks.append(DownBlock2d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)),
|
|
min(max_features, block_expansion * (2 ** (i + 1))),
|
|
kernel_size=3, padding=1))
|
|
self.down_blocks = nn.ModuleList(down_blocks)
|
|
|
|
def forward(self, x):
|
|
outs = [x]
|
|
for down_block in self.down_blocks:
|
|
outs.append(down_block(outs[-1]))
|
|
return outs
|
|
|
|
|
|
class Decoder(nn.Module):
|
|
"""
|
|
Hourglass Decoder
|
|
"""
|
|
|
|
def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
|
|
super(Decoder, self).__init__()
|
|
|
|
up_blocks = []
|
|
|
|
for i in range(num_blocks)[::-1]:
|
|
in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1)))
|
|
out_filters = min(max_features, block_expansion * (2 ** i))
|
|
up_blocks.append(UpBlock2d(in_filters, out_filters, kernel_size=3, padding=1))
|
|
|
|
self.up_blocks = nn.ModuleList(up_blocks)
|
|
self.out_filters = block_expansion + in_features
|
|
|
|
def forward(self, x):
|
|
out = x.pop()
|
|
for up_block in self.up_blocks:
|
|
out = up_block(out)
|
|
skip = x.pop()
|
|
out = torch.cat([out, skip], dim=1)
|
|
return out
|
|
|
|
|
|
class Hourglass(nn.Module):
|
|
"""
|
|
Hourglass architecture.
|
|
"""
|
|
|
|
def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
|
|
super(Hourglass, self).__init__()
|
|
self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features)
|
|
self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features)
|
|
self.out_filters = self.decoder.out_filters
|
|
|
|
def forward(self, x):
|
|
return self.decoder(self.encoder(x))
|
|
|
|
|
|
class AntiAliasInterpolation2d(nn.Module):
|
|
"""
|
|
Band-limited downsampling, for better preservation of the input signal.
|
|
"""
|
|
def __init__(self, channels, scale):
|
|
super(AntiAliasInterpolation2d, self).__init__()
|
|
sigma = (1 / scale - 1) / 2
|
|
kernel_size = 2 * round(sigma * 4) + 1
|
|
self.ka = kernel_size // 2
|
|
self.kb = self.ka - 1 if kernel_size % 2 == 0 else self.ka
|
|
|
|
kernel_size = [kernel_size, kernel_size]
|
|
sigma = [sigma, sigma]
|
|
# The gaussian kernel is the product of the
|
|
# gaussian function of each dimension.
|
|
kernel = 1
|
|
meshgrids = torch.meshgrid(
|
|
[
|
|
torch.arange(size, dtype=torch.float32)
|
|
for size in kernel_size
|
|
]
|
|
)
|
|
for size, std, mgrid in zip(kernel_size, sigma, meshgrids):
|
|
mean = (size - 1) / 2
|
|
kernel *= torch.exp(-(mgrid - mean) ** 2 / (2 * std ** 2))
|
|
|
|
# Make sure sum of values in gaussian kernel equals 1.
|
|
kernel = kernel / torch.sum(kernel)
|
|
# Reshape to depthwise convolutional weight
|
|
kernel = kernel.view(1, 1, *kernel.size())
|
|
kernel = kernel.repeat(channels, *[1] * (kernel.dim() - 1))
|
|
|
|
self.register_buffer('weight', kernel)
|
|
self.groups = channels
|
|
self.scale = scale
|
|
inv_scale = 1 / scale
|
|
self.int_inv_scale = int(inv_scale)
|
|
|
|
def forward(self, input):
|
|
if self.scale == 1.0:
|
|
return input
|
|
|
|
out = F.pad(input, (self.ka, self.kb, self.ka, self.kb))
|
|
out = F.conv2d(out, weight=self.weight, groups=self.groups)
|
|
out = out[:, :, ::self.int_inv_scale, ::self.int_inv_scale]
|
|
|
|
return out
|