| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from models.model_blocks import AdaInResBlock |
| from models.model_blocks import ResBlock |
| from models.model_blocks import UpSamplingBlock |
|
|
|
|
| class SemanticFaceFusionModule(nn.Module): |
| def __init__(self): |
| """ |
| Semantic Face Fusion Module |
| to preserve lighting and background |
| """ |
| super(SemanticFaceFusionModule, self).__init__() |
|
|
| self.sigma = ResBlock(256, 256) |
| self.low_mask_predict = nn.Sequential(nn.Conv2d(256, 1, 3, 1, 1), nn.Sigmoid()) |
| self.z_fuse_block_1 = AdaInResBlock(256, 256) |
| self.z_fuse_block_2 = AdaInResBlock(256, 256) |
|
|
| self.i_low_block = nn.Sequential(nn.LeakyReLU(0.2, inplace=True), nn.Conv2d(256, 3, 3, 1, 1)) |
|
|
| self.f_up = UpSamplingBlock() |
|
|
| def forward(self, target_image, z_enc, z_dec, v_sid): |
| """ |
| Parameters: |
| ---------- |
| target_image: 目标脸图片 |
| z_enc: 1/4原图大小的low-level encoder feature map |
| z_dec: 1/4原图大小的low-level decoder feature map |
| v_sid: the 3D shape aware identity vector |
| |
| Returns: |
| -------- |
| i_r: re-target image |
| i_low: 1/4 size retarget image |
| m_r: face mask |
| m_low: 1/4 size face mask |
| """ |
| z_enc = self.sigma(z_enc) |
|
|
| |
| m_low = self.low_mask_predict(z_dec) |
|
|
| |
| |
| z_fuse = m_low * z_dec + (1 - m_low) * z_enc |
|
|
| z_fuse = self.z_fuse_block_1(z_fuse, v_sid) |
| z_fuse = self.z_fuse_block_2(z_fuse, v_sid) |
|
|
| i_low = self.i_low_block(z_fuse) |
|
|
| i_low = m_low * i_low + (1 - m_low) * F.interpolate(target_image, scale_factor=0.25) |
|
|
| i_r, m_r = self.f_up(z_fuse) |
| i_r = m_r * i_r + (1 - m_r) * target_image |
|
|
| return i_r, i_low, m_r, m_low |
|
|
|
|
| if __name__ == "__main__": |
| import torch |
|
|
| timg = torch.randn(1, 3, 256, 256) |
| z_enc = torch.randn(1, 256, 64, 64) |
| z_dec = torch.randn(1, 256, 64, 64) |
| v_sid = torch.randn(1, 769) |
| model = SemanticFaceFusionModule() |
| i_r, i_low, m_r, m_low = model(timg, z_enc, z_dec, v_sid) |
| print(i_r.shape, i_low.shape, m_r.shape, m_low.shape) |
|
|
