214 lines
7.2 KiB
Python
214 lines
7.2 KiB
Python
# from collections import OrderedDict
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
# from torch.jit.annotations import Dict
|
|
# from torch.nn import BatchNorm2d
|
|
from torchvision import models
|
|
|
|
# from torchvision.models._utils import IntermediateLayerGetter
|
|
from .util import AntiAliasInterpolation2d, DownBlock2d, Hourglass, make_coordinate_grid
|
|
|
|
|
|
class KPDetector(nn.Module):
|
|
"""
|
|
Detecting a keypoints. Return keypoint position and jacobian near each keypoint.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
block_expansion=32,
|
|
num_kp=15,
|
|
num_channels=3,
|
|
max_features=1024,
|
|
num_blocks=5,
|
|
temperature=0.1,
|
|
scale_factor=0.25,
|
|
estimate_jacobian=False,
|
|
estimate_occlusion=False,
|
|
):
|
|
super(KPDetector, self).__init__()
|
|
|
|
self.predictor = Hourglass(
|
|
block_expansion,
|
|
in_features=num_channels,
|
|
max_features=max_features,
|
|
num_blocks=num_blocks,
|
|
)
|
|
|
|
self.kp = nn.Conv2d(
|
|
in_channels=self.predictor.out_filters,
|
|
out_channels=num_kp,
|
|
kernel_size=(7, 7),
|
|
padding=0,
|
|
)
|
|
self.num_kp = num_kp
|
|
|
|
self.estimate_jacobian = estimate_jacobian
|
|
if estimate_jacobian:
|
|
self.num_jacobian_maps = 1
|
|
self.jacobian = nn.Conv2d(
|
|
in_channels=self.predictor.out_filters,
|
|
out_channels=4 * self.num_jacobian_maps,
|
|
kernel_size=(7, 7),
|
|
padding=0,
|
|
)
|
|
self.jacobian.weight.data.zero_()
|
|
self.jacobian.bias.data.copy_(
|
|
torch.tensor([1, 0, 0, 1] * self.num_jacobian_maps, dtype=torch.float)
|
|
)
|
|
|
|
self.estimate_occlusion = estimate_occlusion
|
|
if estimate_occlusion:
|
|
kp_occlusion = [
|
|
DownBlock2d(
|
|
self.predictor.out_filters,
|
|
block_expansion,
|
|
kernel_size=3,
|
|
padding=1,
|
|
)
|
|
]
|
|
kp_occlusion.append(
|
|
DownBlock2d(
|
|
block_expansion, block_expansion * 2, kernel_size=3, padding=1
|
|
)
|
|
)
|
|
kp_occlusion.append(
|
|
DownBlock2d(
|
|
block_expansion * 2, block_expansion * 3, kernel_size=3, padding=1
|
|
)
|
|
)
|
|
kp_occlusion.append(
|
|
DownBlock2d(
|
|
block_expansion * 3, block_expansion * 4, kernel_size=3, padding=1
|
|
)
|
|
)
|
|
kp_occlusion.append(
|
|
nn.Conv2d(
|
|
block_expansion * 4, num_kp, kernel_size=(4, 4), padding=0, stride=4
|
|
)
|
|
)
|
|
self.kp_occlusion = nn.Sequential(*kp_occlusion)
|
|
|
|
self.temperature = temperature
|
|
self.scale_factor = scale_factor
|
|
if self.scale_factor != 1:
|
|
self.down = AntiAliasInterpolation2d(num_channels, self.scale_factor)
|
|
|
|
def create_part_flows(self, source_image, kp_driving, kp_source, bg_params=None):
|
|
"""
|
|
Eq 4. in the paper T_{s<-d}(z)
|
|
"""
|
|
bs, _, h, w = source_image.shape
|
|
num_kp = kp_driving["kp"].shape[1]
|
|
identity_grid = make_coordinate_grid((h, w), type=kp_source["kp"].type())
|
|
identity_grid = identity_grid.view(1, 1, h, w, 2)
|
|
coordinate_grid = identity_grid - kp_driving["kp"].view(bs, num_kp, 1, 1, 2)
|
|
if "jacobian" in kp_driving:
|
|
jacobian = torch.matmul(
|
|
kp_source["jacobian"], torch.inverse(kp_driving["jacobian"])
|
|
)
|
|
jacobian = jacobian.unsqueeze(-3).unsqueeze(-3)
|
|
jacobian = jacobian.repeat(1, 1, h, w, 1, 1)
|
|
coordinate_grid = torch.matmul(jacobian, coordinate_grid.unsqueeze(-1))
|
|
coordinate_grid = coordinate_grid.squeeze(-1)
|
|
# n * 10 * h * w * 2
|
|
driving_to_source = coordinate_grid + kp_source["kp"].view(bs, num_kp, 1, 1, 2)
|
|
flow = driving_to_source - identity_grid
|
|
flow = (h - 1) * flow / 2.0
|
|
return flow.view(bs, num_kp * 2, h, w)
|
|
|
|
# # #adding background feature
|
|
# # identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
|
|
|
|
# # adding background feature
|
|
# if bg_params is None:
|
|
# bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
|
|
# else:
|
|
# bg_grid = identity_grid.repeat(bs, 1, 1, 1, 1)
|
|
# bg_grid = to_homogeneous(bg_grid)
|
|
# bg_grid = torch.matmul(bg_params.view(bs, 1, 1, 1, 3, 3), bg_grid.unsqueeze(-1)).squeeze(-1)
|
|
# bg_grid = from_homogeneous(bg_grid)
|
|
# sparse_motions = torch.cat([bg_grid, driving_to_source], dim=1)
|
|
# return sparse_motions
|
|
|
|
def gaussian2kp(self, heatmap):
|
|
"""
|
|
Extract the mean and from a heatmap
|
|
"""
|
|
shape = heatmap.shape
|
|
heatmap = heatmap.unsqueeze(-1)
|
|
grid = (
|
|
make_coordinate_grid(shape[2:], heatmap.type()).unsqueeze_(0).unsqueeze_(0)
|
|
)
|
|
kp = (heatmap * grid).sum(dim=(2, 3)) # N * 10 * 2
|
|
# kp = {'value': value}
|
|
|
|
return kp
|
|
|
|
def forward(self, x):
|
|
if self.scale_factor != 1:
|
|
x = self.down(x)
|
|
feature_map = self.predictor(x)
|
|
prediction = self.kp(feature_map)
|
|
final_shape = prediction.shape
|
|
heatmap = prediction.view(final_shape[0], final_shape[1], -1)
|
|
heatmap = F.softmax(heatmap / self.temperature, dim=2)
|
|
heatmap = heatmap.view(*final_shape)
|
|
out = self.gaussian2kp(heatmap)
|
|
if self.estimate_jacobian:
|
|
jacobian_map = self.jacobian(feature_map)
|
|
jacobian_map = jacobian_map.reshape(
|
|
final_shape[0],
|
|
self.num_jacobian_maps,
|
|
4,
|
|
final_shape[2],
|
|
final_shape[3],
|
|
)
|
|
heatmap = heatmap.unsqueeze(2)
|
|
jacobian = heatmap * jacobian_map
|
|
jacobian = jacobian.view(final_shape[0], final_shape[1], 4, -1)
|
|
jacobian = jacobian.sum(dim=-1)
|
|
jacobian = jacobian.view(jacobian.shape[0], jacobian.shape[1], 2, 2)
|
|
# N * 10 *2 *2
|
|
out = {"kp": out, "jacobian": jacobian}
|
|
else:
|
|
out = {"kp": out}
|
|
|
|
if self.estimate_occlusion:
|
|
# N*10*1*1
|
|
kp_occlusion = self.kp_occlusion(feature_map)
|
|
kp_occlusion = torch.sigmoid(kp_occlusion)
|
|
out["kp_occlusion"] = kp_occlusion
|
|
# out['heatmap'] = heatmap
|
|
# out['feature_map'] = feature_map
|
|
return out
|
|
|
|
|
|
class TPSKPDetector(nn.Module):
|
|
"""
|
|
Predict K*5 keypoints.
|
|
"""
|
|
|
|
def __init__(self, num_tps, **kwargs):
|
|
super(TPSKPDetector, self).__init__()
|
|
self.num_tps = num_tps
|
|
|
|
self.fg_encoder = models.resnet18(pretrained=False)
|
|
num_features = self.fg_encoder.fc.in_features
|
|
self.fg_encoder.fc = nn.Linear(num_features, num_tps * 5 * 2)
|
|
|
|
def forward(self, image):
|
|
fg_kp = self.fg_encoder(image)
|
|
(
|
|
bs,
|
|
_,
|
|
) = fg_kp.shape
|
|
fg_kp = torch.sigmoid(fg_kp)
|
|
fg_kp = fg_kp * 2 - 1
|
|
out = {"kp": fg_kp.view(bs, self.num_tps * 5, -1)}
|
|
|
|
return out
|