295 lines
11 KiB
Python
295 lines
11 KiB
Python
# from __future__ import absolute_import, division, print_function
|
|
|
|
# import logging
|
|
|
|
# import timm
|
|
# import math
|
|
# import os
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch.autograd import grad
|
|
from torchvision import models
|
|
|
|
from .bg_motion_predictor import BGMotionPredictor
|
|
from .dense_motion import DenseMotionNetwork, TPSDenseMotionNetwork
|
|
|
|
# from .generator import OcclusionAwareGenerator
|
|
from .kp_detector import KPDetector, TPSKPDetector
|
|
from .raft import RaftFlow
|
|
from .transformer.pose_tokenpose_b import get_pose_net
|
|
from .util import AntiAliasInterpolation2d, make_coordinate_grid
|
|
|
|
|
|
class Transform:
|
|
"""
|
|
Random tps transformation for equivariance constraints. See Sec 3.3
|
|
"""
|
|
|
|
def __init__(self, bs, **kwargs):
|
|
noise = torch.normal(
|
|
mean=0, std=kwargs["sigma_affine"] * torch.ones([bs, 2, 3])
|
|
)
|
|
self.theta = noise + torch.eye(2, 3).view(1, 2, 3)
|
|
self.bs = bs
|
|
|
|
if ("sigma_tps" in kwargs) and ("points_tps" in kwargs):
|
|
self.tps = True
|
|
self.control_points = make_coordinate_grid(
|
|
(kwargs["points_tps"], kwargs["points_tps"]), type=noise.type()
|
|
)
|
|
self.control_points = self.control_points.unsqueeze(0)
|
|
self.control_params = torch.normal(
|
|
mean=0,
|
|
std=kwargs["sigma_tps"]
|
|
* torch.ones([bs, 1, kwargs["points_tps"] ** 2]),
|
|
)
|
|
else:
|
|
self.tps = False
|
|
|
|
def transform_frame(self, frame):
|
|
grid = make_coordinate_grid(frame.shape[2:], type=frame.type()).unsqueeze(0)
|
|
grid = grid.view(1, frame.shape[2] * frame.shape[3], 2)
|
|
grid = self.warp_coordinates(grid).view(
|
|
self.bs, frame.shape[2], frame.shape[3], 2
|
|
)
|
|
return F.grid_sample(frame, grid, padding_mode="reflection")
|
|
|
|
def warp_coordinates(self, coordinates):
|
|
theta = self.theta.type(coordinates.type())
|
|
theta = theta.unsqueeze(1)
|
|
# N * (HW) * 2 * 1
|
|
transformed = (
|
|
torch.matmul(theta[:, :, :, :2], coordinates.unsqueeze(-1))
|
|
+ theta[:, :, :, 2:]
|
|
)
|
|
transformed = transformed.squeeze(-1)
|
|
|
|
if self.tps:
|
|
control_points = self.control_points.type(coordinates.type())
|
|
control_params = self.control_params.type(coordinates.type())
|
|
# 1 * (HW) * 1 * 2 - 1 * 1 * 25 *2
|
|
distances = coordinates.view(
|
|
coordinates.shape[0], -1, 1, 2
|
|
) - control_points.view(1, 1, -1, 2)
|
|
distances = torch.abs(distances).sum(-1)
|
|
|
|
result = distances**2
|
|
result = result * torch.log(distances + 1e-6)
|
|
result = result * control_params
|
|
result = result.sum(dim=2).view(self.bs, coordinates.shape[1], 1)
|
|
transformed = transformed + result
|
|
|
|
return transformed
|
|
|
|
def jacobian(self, coordinates):
|
|
new_coordinates = self.warp_coordinates(coordinates)
|
|
grad_x = grad(new_coordinates[..., 0].sum(), coordinates, create_graph=True)
|
|
grad_y = grad(new_coordinates[..., 1].sum(), coordinates, create_graph=True)
|
|
jacobian = torch.cat([grad_x[0].unsqueeze(-2), grad_y[0].unsqueeze(-2)], dim=-2)
|
|
return jacobian
|
|
|
|
|
|
class Vgg19(torch.nn.Module):
|
|
"""
|
|
Vgg19 network for perceptual loss. See Sec 3.3.
|
|
"""
|
|
|
|
def __init__(self, requires_grad=False):
|
|
super(Vgg19, self).__init__()
|
|
vgg_pretrained_features = models.vgg19(pretrained=True).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])
|
|
|
|
self.mean = torch.nn.Parameter(
|
|
data=torch.Tensor(np.array([0.485, 0.456, 0.406]).reshape((1, 3, 1, 1))),
|
|
requires_grad=False,
|
|
)
|
|
self.std = torch.nn.Parameter(
|
|
data=torch.Tensor(np.array([0.229, 0.224, 0.225]).reshape((1, 3, 1, 1))),
|
|
requires_grad=False,
|
|
)
|
|
|
|
if not requires_grad:
|
|
for param in self.parameters():
|
|
param.requires_grad = False
|
|
|
|
def forward(self, X):
|
|
# X = X.clamp(-1, 1)
|
|
# X = X / 2 + 0.5
|
|
X = (X - self.mean) / self.std
|
|
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
|
|
|
|
|
|
class ImagePyramide(torch.nn.Module):
|
|
"""
|
|
Create image pyramide for computing pyramide perceptual loss. See Sec 3.3
|
|
"""
|
|
|
|
def __init__(self, scales, num_channels=3):
|
|
super(ImagePyramide, self).__init__()
|
|
downs = {}
|
|
for scale in scales:
|
|
downs[str(scale).replace(".", "-")] = AntiAliasInterpolation2d(
|
|
num_channels, scale
|
|
)
|
|
self.downs = nn.ModuleDict(downs)
|
|
|
|
def forward(self, x):
|
|
out_dict = {}
|
|
for scale, down_module in self.downs.items():
|
|
out_dict["prediction_" + str(scale).replace("-", ".")] = down_module(x)
|
|
return out_dict
|
|
|
|
|
|
class MRFA(nn.Module):
|
|
def __init__(self, cfg, **kwargs):
|
|
super(MRFA, self).__init__()
|
|
|
|
self.cfg = cfg
|
|
train_params = cfg["train_params"]
|
|
self.train_params = train_params
|
|
self.scales = train_params["scales"]
|
|
self.loss_weights = train_params["loss_weights"]
|
|
self.pyramid = ImagePyramide(self.scales).cuda()
|
|
if sum(self.loss_weights["perceptual"]) != 0:
|
|
self.vgg = Vgg19().cuda()
|
|
|
|
##################################################
|
|
if train_params["prior_model"] == "fomm":
|
|
self.encoder = KPDetector(**cfg["fomm_kp_detector"])
|
|
self.dense_motion = DenseMotionNetwork(**cfg["dense_motion"])
|
|
elif train_params["prior_model"] == "tpsm":
|
|
self.encoder = TPSKPDetector(**cfg["tpsm_kp_detector"])
|
|
self.dense_motion = TPSDenseMotionNetwork(**cfg["tpsm_dense_motion"])
|
|
self.dropout_epoch = train_params["dropout_epoch"]
|
|
self.dropout_maxp = train_params["dropout_maxp"]
|
|
self.dropout_startp = train_params["dropout_startp"]
|
|
self.dropout_inc_epoch = train_params["dropout_inc_epoch"]
|
|
elif train_params["prior_model"] == "mtia":
|
|
self.encoder = get_pose_net(cfg["mtia_kp_detector"], is_train=True)
|
|
self.dense_motion = DenseMotionNetwork(**cfg["dense_motion"])
|
|
else:
|
|
print("please specify the prior motion model")
|
|
|
|
if train_params["bg_start"] < train_params["num_epochs"]:
|
|
self.bg_predictor = BGMotionPredictor()
|
|
self.bg_start = train_params["bg_start"]
|
|
|
|
self.decoder = RaftFlow(**cfg["raft_flow"])
|
|
self.down = AntiAliasInterpolation2d(3, 0.25)
|
|
|
|
def forward(self, x, epoch=100, is_train=True):
|
|
kp_s = self.encoder(x["source"])
|
|
kp_d = self.encoder(x["driving"])
|
|
|
|
img_down = self.down(x["source"])
|
|
if epoch >= self.bg_start:
|
|
bg_param = self.bg_predictor(x["source"], x["driving"])
|
|
else:
|
|
bg_param = None
|
|
|
|
if self.train_params["prior_model"] == "tpsm" and epoch < self.dropout_epoch:
|
|
# dropout_p will linearly increase from dropout_startp to dropout_maxp
|
|
dropout_flag = True
|
|
dropout_p = min(
|
|
epoch / self.dropout_inc_epoch * self.dropout_maxp
|
|
+ self.dropout_startp,
|
|
self.dropout_maxp,
|
|
)
|
|
else:
|
|
dropout_flag = False
|
|
dropout_p = 0
|
|
dense_motion = self.dense_motion(
|
|
x["source"],
|
|
kp_d,
|
|
kp_s,
|
|
bg_param=bg_param,
|
|
dropout_flag=dropout_flag,
|
|
dropout_p=dropout_p,
|
|
)
|
|
|
|
if self.train_params["prior_model"] == "tpsm":
|
|
kp_s_value = kp_s["kp"].view(x["source"].shape[0], -1, 5, 2).mean(2)
|
|
kp_d_value = kp_d["kp"].view(x["driving"].shape[0], -1, 5, 2).mean(2)
|
|
else:
|
|
kp_s_value = kp_s["kp"]
|
|
kp_d_value = kp_d["kp"]
|
|
|
|
gen, warp_img, occlusion = self.decoder(
|
|
kp_s_value, kp_d_value, dense_motion, img=img_down, img_full=x["source"]
|
|
)
|
|
warp_img = torch.cat([warp_img, occlusion.repeat(1, 3, 1, 1)], dim=3)
|
|
|
|
loss_values = {}
|
|
|
|
if not is_train:
|
|
return gen, warp_img, loss_values, kp_s["kp"], kp_d["kp"]
|
|
|
|
pyramide_real = self.pyramid(x["driving"])
|
|
pyramide_generated = self.pyramid(gen)
|
|
value_total = 0
|
|
for scale in self.scales:
|
|
x_vgg = self.vgg(pyramide_generated["prediction_" + str(scale)])
|
|
y_vgg = self.vgg(pyramide_real["prediction_" + str(scale)])
|
|
|
|
for i, weight in enumerate(self.loss_weights["perceptual"]):
|
|
value = torch.abs(x_vgg[i] - y_vgg[i].detach()).mean()
|
|
value_total += self.loss_weights["perceptual"][i] * value
|
|
loss_values["perceptual"] = value_total
|
|
|
|
if self.loss_weights["equivariance"] != 0:
|
|
transform = Transform(
|
|
x["driving"].shape[0], **self.train_params["transform_params"]
|
|
)
|
|
transformed_frame = transform.transform_frame(x["driving"])
|
|
transformed_kp = self.encoder(transformed_frame)
|
|
value = torch.abs(
|
|
kp_d["kp"] - transform.warp_coordinates(transformed_kp["kp"])
|
|
)
|
|
value = value.mean()
|
|
loss_values["equivariance"] = self.loss_weights["equivariance"] * value
|
|
|
|
if self.loss_weights["equivariance_jacobian"] != 0:
|
|
jacobian_transformed = torch.matmul(
|
|
transform.jacobian(transformed_kp["kp"]), transformed_kp["jacobian"]
|
|
)
|
|
normed_driving = torch.inverse(kp_d["jacobian"])
|
|
normed_transformed = jacobian_transformed
|
|
value = torch.matmul(normed_driving, normed_transformed)
|
|
eye = torch.eye(2).view(1, 1, 2, 2).type(value.type())
|
|
value = torch.abs(eye - value)
|
|
loss_values["equivariance_jacobian"] = (
|
|
self.loss_weights["equivariance_jacobian"] * value
|
|
)
|
|
|
|
if epoch >= self.bg_start:
|
|
bg_param_reverse = self.bg_predictor(x["driving"], x["source"])
|
|
value = torch.matmul(bg_param, bg_param_reverse)
|
|
eye = torch.eye(3).view(1, 1, 3, 3).type(value.type())
|
|
value = torch.abs(eye - value).mean()
|
|
loss_values["bg"] = 10 * value
|
|
|
|
return gen, warp_img, loss_values, kp_s["kp"], kp_d["kp"]
|