127 lines
6.3 KiB
Python
127 lines
6.3 KiB
Python
"""
|
|
Copyright Snap Inc. 2021. This sample code is made available by Snap Inc. for informational purposes only.
|
|
No license, whether implied or otherwise, is granted in or to such code (including any rights to copy, modify,
|
|
publish, distribute and/or commercialize such code), unless you have entered into a separate agreement for such rights.
|
|
Such code is provided as-is, without warranty of any kind, express or implied, including any warranties of merchantability,
|
|
title, fitness for a particular purpose, non-infringement, or that such code is free of defects, errors or viruses.
|
|
In no event will Snap Inc. be liable for any damages or losses of any kind arising from the sample code or your use thereof.
|
|
"""
|
|
|
|
import torch
|
|
from torch import nn
|
|
import torch.nn.functional as F
|
|
from .util import ResBlock2d, SameBlock2d, UpBlock2d, DownBlock2d
|
|
from .pixelwise_flow_predictor import PixelwiseFlowPredictor
|
|
|
|
|
|
class Generator(nn.Module):
|
|
"""
|
|
Generator that given source image and region parameters try to transform image according to movement trajectories
|
|
induced by region parameters. Generator follows Johnson architecture.
|
|
"""
|
|
|
|
def __init__(self, num_channels, num_regions, block_expansion, max_features, num_down_blocks,
|
|
num_bottleneck_blocks, pixelwise_flow_predictor_params=None, skips=False, revert_axis_swap=True):
|
|
super(Generator, self).__init__()
|
|
|
|
if pixelwise_flow_predictor_params is not None:
|
|
self.pixelwise_flow_predictor = PixelwiseFlowPredictor(num_regions=num_regions, num_channels=num_channels,
|
|
revert_axis_swap=revert_axis_swap,
|
|
**pixelwise_flow_predictor_params)
|
|
else:
|
|
self.pixelwise_flow_predictor = None
|
|
|
|
self.first = SameBlock2d(num_channels, block_expansion, kernel_size=(7, 7), padding=(3, 3))
|
|
|
|
down_blocks = []
|
|
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)))
|
|
self.down_blocks = nn.ModuleList(down_blocks)
|
|
|
|
up_blocks = []
|
|
for i in range(num_down_blocks):
|
|
in_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i)))
|
|
out_features = min(max_features, block_expansion * (2 ** (num_down_blocks - i - 1)))
|
|
up_blocks.append(UpBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
|
|
self.up_blocks = nn.ModuleList(up_blocks)
|
|
|
|
self.bottleneck = torch.nn.Sequential()
|
|
in_features = min(max_features, block_expansion * (2 ** num_down_blocks))
|
|
for i in range(num_bottleneck_blocks):
|
|
self.bottleneck.add_module('r' + str(i), ResBlock2d(in_features, kernel_size=(3, 3), padding=(1, 1)))
|
|
|
|
self.final = nn.Conv2d(block_expansion, num_channels, kernel_size=(7, 7), padding=(3, 3))
|
|
self.num_channels = num_channels
|
|
self.skips = skips
|
|
|
|
@staticmethod
|
|
def deform_input(inp, optical_flow):
|
|
_, h_old, w_old, _ = optical_flow.shape
|
|
_, _, h, w = inp.shape
|
|
if h_old != h or w_old != w:
|
|
optical_flow = optical_flow.permute(0, 3, 1, 2)
|
|
optical_flow = F.interpolate(optical_flow, size=(h, w), mode='bilinear')
|
|
optical_flow = optical_flow.permute(0, 2, 3, 1)
|
|
return F.grid_sample(inp, optical_flow)
|
|
|
|
def apply_optical(self, input_previous=None, input_skip=None, motion_params=None):
|
|
if motion_params is not None:
|
|
if 'occlusion_map' in motion_params:
|
|
occlusion_map = motion_params['occlusion_map']
|
|
else:
|
|
occlusion_map = None
|
|
deformation = motion_params['optical_flow']
|
|
input_skip = self.deform_input(input_skip, deformation)
|
|
|
|
if occlusion_map is not None:
|
|
if input_skip.shape[2] != occlusion_map.shape[2] or input_skip.shape[3] != occlusion_map.shape[3]:
|
|
occlusion_map = F.interpolate(occlusion_map, size=input_skip.shape[2:], mode='bilinear')
|
|
if input_previous is not None:
|
|
input_skip = input_skip * occlusion_map + input_previous * (1 - occlusion_map)
|
|
else:
|
|
input_skip = input_skip * occlusion_map
|
|
out = input_skip
|
|
else:
|
|
out = input_previous if input_previous is not None else input_skip
|
|
return out
|
|
|
|
def forward(self, source_image, driving_region_params, source_region_params, bg_params=None):
|
|
out = self.first(source_image)
|
|
skips = [out]
|
|
for i in range(len(self.down_blocks)):
|
|
out = self.down_blocks[i](out)
|
|
skips.append(out)
|
|
|
|
output_dict = {}
|
|
if self.pixelwise_flow_predictor is not None:
|
|
motion_params = self.pixelwise_flow_predictor(source_image=source_image,
|
|
driving_region_params=driving_region_params,
|
|
source_region_params=source_region_params,
|
|
bg_params=bg_params)
|
|
output_dict["deformed"] = self.deform_input(source_image, motion_params['optical_flow'])
|
|
if 'occlusion_map' in motion_params:
|
|
output_dict['occlusion_map'] = motion_params['occlusion_map']
|
|
else:
|
|
motion_params = None
|
|
|
|
out = self.apply_optical(input_previous=None, input_skip=out, motion_params=motion_params)
|
|
|
|
out = self.bottleneck(out)
|
|
for i in range(len(self.up_blocks)):
|
|
if self.skips:
|
|
out = self.apply_optical(input_skip=skips[-(i + 1)], input_previous=out, motion_params=motion_params)
|
|
out = self.up_blocks[i](out)
|
|
if self.skips:
|
|
out = self.apply_optical(input_skip=skips[0], input_previous=out, motion_params=motion_params)
|
|
out = self.final(out)
|
|
out = F.sigmoid(out)
|
|
|
|
if self.skips:
|
|
out = self.apply_optical(input_skip=source_image, input_previous=out, motion_params=motion_params)
|
|
|
|
output_dict["prediction"] = out
|
|
|
|
return output_dict
|