165 lines
5.7 KiB
Python
165 lines
5.7 KiB
Python
import torch
|
|
import torch.nn.functional as F
|
|
from torch import nn
|
|
|
|
from .util import DownBlock2d, ResBlock2d, SameBlock2d, UpBlock2d
|
|
|
|
|
|
class InpaintingNetwork(nn.Module):
|
|
"""
|
|
Inpaint the missing regions and reconstruct the Driving image.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
num_channels,
|
|
block_expansion,
|
|
max_features,
|
|
num_down_blocks,
|
|
multi_mask=True,
|
|
**kwargs,
|
|
):
|
|
super(InpaintingNetwork, self).__init__()
|
|
|
|
self.num_down_blocks = num_down_blocks
|
|
self.multi_mask = multi_mask
|
|
self.first = SameBlock2d(
|
|
num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3)
|
|
)
|
|
|
|
down_blocks = []
|
|
up_blocks = []
|
|
resblock = []
|
|
for i in range(num_down_blocks):
|
|
in_features = min(max_features, block_expansion * (2**i))
|
|
out_features = min(max_features, block_expansion * (2 ** (i + 1)))
|
|
down_blocks.append(
|
|
DownBlock2d(
|
|
in_features, out_features, kernel_size=(3, 3), padding=(1, 1)
|
|
)
|
|
)
|
|
decoder_in_feature = out_features * 2
|
|
if i == num_down_blocks - 1:
|
|
decoder_in_feature = out_features
|
|
up_blocks.append(
|
|
UpBlock2d(
|
|
decoder_in_feature, in_features, kernel_size=(3, 3), padding=(1, 1)
|
|
)
|
|
)
|
|
resblock.append(
|
|
ResBlock2d(decoder_in_feature, kernel_size=(3, 3), padding=(1, 1))
|
|
)
|
|
resblock.append(
|
|
ResBlock2d(decoder_in_feature, kernel_size=(3, 3), padding=(1, 1))
|
|
)
|
|
self.down_blocks = nn.ModuleList(down_blocks)
|
|
self.up_blocks = nn.ModuleList(up_blocks[::-1])
|
|
self.resblock = nn.ModuleList(resblock[::-1])
|
|
|
|
self.final = nn.Conv2d(
|
|
block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3)
|
|
)
|
|
self.num_channels = num_channels
|
|
|
|
def deform_input(self, inp, deformation):
|
|
_, h_old, w_old, _ = deformation.shape
|
|
_, _, h, w = inp.shape
|
|
if h_old != h or w_old != w:
|
|
deformation = deformation.permute(0, 3, 1, 2)
|
|
deformation = F.interpolate(
|
|
deformation, size=(h, w), mode="bilinear", align_corners=True
|
|
)
|
|
deformation = deformation.permute(0, 2, 3, 1)
|
|
return F.grid_sample(inp, deformation, align_corners=True)
|
|
|
|
def occlude_input(self, inp, occlusion_map):
|
|
if not self.multi_mask:
|
|
if (
|
|
inp.shape[2] != occlusion_map.shape[2]
|
|
or inp.shape[3] != occlusion_map.shape[3]
|
|
):
|
|
occlusion_map = F.interpolate(
|
|
occlusion_map,
|
|
size=inp.shape[2:],
|
|
mode="bilinear",
|
|
align_corners=True,
|
|
)
|
|
out = inp * occlusion_map
|
|
return out
|
|
|
|
def forward(self, source_image, dense_motion):
|
|
out = self.first(source_image)
|
|
encoder_map = [out]
|
|
for i in range(len(self.down_blocks)):
|
|
out = self.down_blocks[i](out)
|
|
encoder_map.append(out)
|
|
|
|
output_dict = {}
|
|
output_dict["contribution_maps"] = dense_motion["contribution_maps"]
|
|
output_dict["deformed_source"] = dense_motion["deformed_source"]
|
|
|
|
occlusion_map = dense_motion["occlusion_map"]
|
|
output_dict["occlusion_map"] = occlusion_map
|
|
|
|
deformation = dense_motion["deformation"]
|
|
out_ij = self.deform_input(out.detach(), deformation)
|
|
out = self.deform_input(out, deformation)
|
|
|
|
out_ij = self.occlude_input(out_ij, occlusion_map[0].detach())
|
|
out = self.occlude_input(out, occlusion_map[0])
|
|
|
|
warped_encoder_maps = []
|
|
warped_encoder_maps.append(out_ij)
|
|
|
|
for i in range(self.num_down_blocks):
|
|
out = self.resblock[2 * i](out)
|
|
out = self.resblock[2 * i + 1](out)
|
|
out = self.up_blocks[i](out)
|
|
|
|
encode_i = encoder_map[-(i + 2)]
|
|
encode_ij = self.deform_input(encode_i.detach(), deformation)
|
|
encode_i = self.deform_input(encode_i, deformation)
|
|
|
|
occlusion_ind = 0
|
|
if self.multi_mask:
|
|
occlusion_ind = i + 1
|
|
encode_ij = self.occlude_input(
|
|
encode_ij, occlusion_map[occlusion_ind].detach()
|
|
)
|
|
encode_i = self.occlude_input(encode_i, occlusion_map[occlusion_ind])
|
|
warped_encoder_maps.append(encode_ij)
|
|
|
|
if i == self.num_down_blocks - 1:
|
|
break
|
|
|
|
out = torch.cat([out, encode_i], 1)
|
|
|
|
deformed_source = self.deform_input(source_image, deformation)
|
|
output_dict["deformed"] = deformed_source
|
|
output_dict["warped_encoder_maps"] = warped_encoder_maps
|
|
|
|
occlusion_last = occlusion_map[-1]
|
|
if not self.multi_mask:
|
|
occlusion_last = F.interpolate(
|
|
occlusion_last, size=out.shape[2:], mode="bilinear", align_corners=True
|
|
)
|
|
|
|
out = out * (1 - occlusion_last) + encode_i
|
|
out = self.final(out)
|
|
out = torch.sigmoid(out)
|
|
out = out * (1 - occlusion_last) + deformed_source * occlusion_last
|
|
output_dict["prediction"] = out
|
|
|
|
return output_dict
|
|
|
|
def get_encode(self, driver_image, occlusion_map):
|
|
out = self.first(driver_image)
|
|
encoder_map = []
|
|
encoder_map.append(self.occlude_input(out.detach(), occlusion_map[-1].detach()))
|
|
for i in range(len(self.down_blocks)):
|
|
out = self.down_blocks[i](out.detach())
|
|
out_mask = self.occlude_input(out.detach(), occlusion_map[2 - i].detach())
|
|
encoder_map.append(out_mask.detach())
|
|
|
|
return encoder_map
|