-
Notifications
You must be signed in to change notification settings - Fork 286
/
two_algo_face_rotator.py
89 lines (77 loc) · 3.62 KB
/
two_algo_face_rotator.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import torch
from torch import Tensor
from torch.nn import Sequential, Tanh, Sigmoid
from torch.nn.functional import affine_grid, grid_sample
from tha.batch_input_module import BatchInputModule, BatchInputModuleSpec
from nn.conv import Conv7
from nn.encoder_decoder_module import EncoderDecoderModule
class TwoAlgoFaceRotator(BatchInputModule):
def __init__(self,
image_size: int = 256,
image_channels: int = 4,
pose_size: int = 3,
intermediate_channels: int = 64,
bottleneck_image_size: int = 32,
bottleneck_block_count: int = 6,
initialization_method: str = 'he'):
super().__init__()
self.main_body = EncoderDecoderModule(
image_size=image_size,
image_channels=image_channels + pose_size,
output_channels=intermediate_channels,
bottleneck_image_size=bottleneck_image_size,
bottleneck_block_count=bottleneck_block_count,
initialization_method=initialization_method)
self.pumarola_color_change = Sequential(
Conv7(intermediate_channels, image_channels, initialization_method),
Tanh())
self.pumarola_alpha_mask = Sequential(
Conv7(intermediate_channels, image_channels, initialization_method),
Sigmoid())
self.zhou_grid_change = Conv7(intermediate_channels, 2, initialization_method)
def forward(self, image: Tensor, pose: Tensor):
n = image.size(0)
c = image.size(1)
h = image.size(2)
w = image.size(3)
pose = pose.unsqueeze(2).unsqueeze(3)
pose = pose.expand(pose.size(0), pose.size(1), image.size(2), image.size(3))
x = torch.cat([image, pose], dim=1)
y = self.main_body(x)
color_change = self.pumarola_color_change(y)
alpha_mask = self.pumarola_alpha_mask(y)
color_changed = alpha_mask * image + (1 - alpha_mask) * color_change
grid_change = torch.transpose(self.zhou_grid_change(y).view(n, 2, h * w), 1, 2).view(n, h, w, 2)
device = self.zhou_grid_change.weight.device
identity = torch.Tensor([[1, 0, 0], [0, 1, 0]]).to(device).unsqueeze(0).repeat(n, 1, 1)
base_grid = affine_grid(identity, [n, c, h, w])
grid = base_grid + grid_change
resampled = grid_sample(image, grid, mode='bilinear', padding_mode='border')
return [color_changed, resampled, color_change, alpha_mask, grid_change, grid]
def forward_from_batch(self, batch):
return self.forward(batch[0], batch[1])
class TwoAlgoFaceRotatorSpec(BatchInputModuleSpec):
def __init__(self,
image_size: int = 256,
image_channels: int = 4,
pose_size: int = 3,
intermediate_channels: int = 64,
bottleneck_image_size: int = 32,
bottleneck_block_count: int = 6,
initialization_method: str = 'he'):
self.image_size = image_size
self.image_channels = image_channels
self.pose_size = pose_size
self.intermediate_channels = intermediate_channels
self.bottleneck_image_size = bottleneck_image_size
self.bottleneck_block_count = bottleneck_block_count
self.initialization_method = initialization_method
def get_module(self) -> BatchInputModule:
return TwoAlgoFaceRotator(
self.image_size,
self.image_channels,
self.pose_size,
self.intermediate_channels,
self.bottleneck_image_size,
self.bottleneck_block_count,
self.initialization_method)