Spaces:
Runtime error
Runtime error
File size: 3,189 Bytes
882f6e2 |
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 |
"""
Copyright (c) Meta Platforms, Inc. and affiliates.
All rights reserved.
This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""
from typing import Dict, Tuple
import numpy as np
import torch as th
import torch.nn as nn
import visualize.ca_body.nn.layers as la
from attrdict import AttrDict
class FaceDecoderFrontal(nn.Module):
def __init__(
self,
assets: AttrDict,
n_latent: int = 256,
n_vert_out: int = 3 * 7306,
tex_out_shp: Tuple[int, int] = (1024, 1024),
tex_roi: Tuple[Tuple[int, int], Tuple[int, int]] = ((0, 0), (1024, 1024)),
) -> None:
super().__init__()
self.n_latent = n_latent
self.n_vert_out = n_vert_out
self.tex_roi = tex_roi
self.tex_roi_shp: Tuple[int, int] = tuple(
[int(i) for i in np.diff(np.array(tex_roi), axis=0).squeeze()]
)
self.tex_out_shp = tex_out_shp
self.encmod = nn.Sequential(
la.LinearWN(n_latent, 256), nn.LeakyReLU(0.2, inplace=True)
)
self.geommod = nn.Sequential(la.LinearWN(256, n_vert_out))
self.viewmod = nn.Sequential(la.LinearWN(3, 8), nn.LeakyReLU(0.2, inplace=True))
self.texmod2 = nn.Sequential(
la.LinearWN(256 + 8, 256 * 4 * 4), nn.LeakyReLU(0.2, inplace=True)
)
self.texmod = nn.Sequential(
la.ConvTranspose2dWNUB(256, 256, 8, 8, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(256, 128, 16, 16, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(128, 128, 32, 32, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(128, 64, 64, 64, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(64, 64, 128, 128, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(64, 32, 256, 256, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(32, 8, 512, 512, 4, 2, 1),
nn.LeakyReLU(0.2, inplace=True),
la.ConvTranspose2dWNUB(8, 3, 1024, 1024, 4, 2, 1),
)
self.bias = nn.Parameter(th.zeros(3, self.tex_roi_shp[0], self.tex_roi_shp[1]))
self.bias.data.zero_()
self.register_buffer(
"frontal_view", th.as_tensor(assets.face_frontal_view, dtype=th.float32)
)
self.apply(lambda x: la.glorot(x, 0.2))
la.glorot(self.texmod[-1], 1.0)
def forward(self, face_embs: th.Tensor) -> Dict[str, th.Tensor]:
B = face_embs.shape[0]
view = self.frontal_view[np.newaxis].expand(B, -1)
encout = self.encmod(face_embs)
geomout = self.geommod(encout)
viewout = self.viewmod(view)
encview = th.cat([encout, viewout], dim=1)
texout = self.texmod(self.texmod2(encview).view(-1, 256, 4, 4))
out = {"face_geom": geomout.view(geomout.shape[0], -1, 3)}
out["face_tex_raw"] = texout
texout = texout + self.bias[None]
out["face_tex"] = 255 * (texout + 0.5)
return out
|