|
import torch |
|
import torch.nn as nn |
|
import numpy as np |
|
import os |
|
from pytorch3d.structures import Meshes |
|
from pytorch3d.renderer import ( |
|
look_at_view_transform, |
|
PerspectiveCameras, |
|
FoVPerspectiveCameras, |
|
PointLights, |
|
DirectionalLights, |
|
Materials, |
|
RasterizationSettings, |
|
MeshRenderer, |
|
MeshRasterizer, |
|
SoftPhongShader, |
|
TexturesUV, |
|
TexturesVertex, |
|
blending, |
|
) |
|
|
|
from pytorch3d.ops import interpolate_face_attributes |
|
|
|
from pytorch3d.renderer.blending import ( |
|
BlendParams, |
|
hard_rgb_blend, |
|
sigmoid_alpha_blend, |
|
softmax_rgb_blend, |
|
) |
|
|
|
|
|
class SoftSimpleShader(nn.Module): |
|
""" |
|
Per pixel lighting - the lighting model is applied using the interpolated |
|
coordinates and normals for each pixel. The blending function returns the |
|
soft aggregated color using all the faces per pixel. |
|
|
|
To use the default values, simply initialize the shader with the desired |
|
device e.g. |
|
|
|
""" |
|
|
|
def __init__( |
|
self, device="cpu", cameras=None, lights=None, materials=None, blend_params=None |
|
): |
|
super().__init__() |
|
self.lights = lights if lights is not None else PointLights(device=device) |
|
self.materials = ( |
|
materials if materials is not None else Materials(device=device) |
|
) |
|
self.cameras = cameras |
|
self.blend_params = blend_params if blend_params is not None else BlendParams() |
|
|
|
def to(self, device): |
|
|
|
self.cameras = self.cameras.to(device) |
|
self.materials = self.materials.to(device) |
|
self.lights = self.lights.to(device) |
|
return self |
|
|
|
def forward(self, fragments, meshes, **kwargs) -> torch.Tensor: |
|
|
|
texels = meshes.sample_textures(fragments) |
|
blend_params = kwargs.get("blend_params", self.blend_params) |
|
|
|
cameras = kwargs.get("cameras", self.cameras) |
|
if cameras is None: |
|
msg = "Cameras must be specified either at initialization \ |
|
or in the forward pass of SoftPhongShader" |
|
raise ValueError(msg) |
|
znear = kwargs.get("znear", getattr(cameras, "znear", 1.0)) |
|
zfar = kwargs.get("zfar", getattr(cameras, "zfar", 100.0)) |
|
images = softmax_rgb_blend( |
|
texels, fragments, blend_params, znear=znear, zfar=zfar |
|
) |
|
return images |
|
|
|
|
|
class Render_3DMM(nn.Module): |
|
def __init__( |
|
self, |
|
focal=1015, |
|
img_h=500, |
|
img_w=500, |
|
batch_size=1, |
|
device=torch.device("cuda:0"), |
|
): |
|
super(Render_3DMM, self).__init__() |
|
|
|
self.focal = focal |
|
self.img_h = img_h |
|
self.img_w = img_w |
|
self.device = device |
|
self.renderer = self.get_render(batch_size) |
|
|
|
dir_path = os.path.dirname(os.path.realpath(__file__)) |
|
topo_info = np.load( |
|
os.path.join(dir_path, "3DMM", "topology_info.npy"), allow_pickle=True |
|
).item() |
|
self.tris = torch.as_tensor(topo_info["tris"]).to(self.device) |
|
self.vert_tris = torch.as_tensor(topo_info["vert_tris"]).to(self.device) |
|
|
|
def compute_normal(self, geometry): |
|
vert_1 = torch.index_select(geometry, 1, self.tris[:, 0]) |
|
vert_2 = torch.index_select(geometry, 1, self.tris[:, 1]) |
|
vert_3 = torch.index_select(geometry, 1, self.tris[:, 2]) |
|
nnorm = torch.cross(vert_2 - vert_1, vert_3 - vert_1, 2) |
|
tri_normal = nn.functional.normalize(nnorm, dim=2) |
|
v_norm = tri_normal[:, self.vert_tris, :].sum(2) |
|
vert_normal = v_norm / v_norm.norm(dim=2).unsqueeze(2) |
|
return vert_normal |
|
|
|
def get_render(self, batch_size=1): |
|
half_s = self.img_w * 0.5 |
|
R, T = look_at_view_transform(10, 0, 0) |
|
R = R.repeat(batch_size, 1, 1) |
|
T = torch.zeros((batch_size, 3), dtype=torch.float32).to(self.device) |
|
|
|
cameras = FoVPerspectiveCameras( |
|
device=self.device, |
|
R=R, |
|
T=T, |
|
znear=0.01, |
|
zfar=20, |
|
fov=2 * np.arctan(self.img_w // 2 / self.focal) * 180.0 / np.pi, |
|
) |
|
lights = PointLights( |
|
device=self.device, |
|
location=[[0.0, 0.0, 1e5]], |
|
ambient_color=[[1, 1, 1]], |
|
specular_color=[[0.0, 0.0, 0.0]], |
|
diffuse_color=[[0.0, 0.0, 0.0]], |
|
) |
|
sigma = 1e-4 |
|
raster_settings = RasterizationSettings( |
|
image_size=(self.img_h, self.img_w), |
|
blur_radius=np.log(1.0 / 1e-4 - 1.0) * sigma / 18.0, |
|
faces_per_pixel=2, |
|
perspective_correct=False, |
|
) |
|
blend_params = blending.BlendParams(background_color=[0, 0, 0]) |
|
renderer = MeshRenderer( |
|
rasterizer=MeshRasterizer(raster_settings=raster_settings, cameras=cameras), |
|
shader=SoftSimpleShader( |
|
lights=lights, blend_params=blend_params, cameras=cameras |
|
), |
|
) |
|
return renderer.to(self.device) |
|
|
|
@staticmethod |
|
def Illumination_layer(face_texture, norm, gamma): |
|
|
|
n_b, num_vertex, _ = face_texture.size() |
|
n_v_full = n_b * num_vertex |
|
gamma = gamma.view(-1, 3, 9).clone() |
|
gamma[:, :, 0] += 0.8 |
|
|
|
gamma = gamma.permute(0, 2, 1) |
|
|
|
a0 = np.pi |
|
a1 = 2 * np.pi / np.sqrt(3.0) |
|
a2 = 2 * np.pi / np.sqrt(8.0) |
|
c0 = 1 / np.sqrt(4 * np.pi) |
|
c1 = np.sqrt(3.0) / np.sqrt(4 * np.pi) |
|
c2 = 3 * np.sqrt(5.0) / np.sqrt(12 * np.pi) |
|
d0 = 0.5 / np.sqrt(3.0) |
|
|
|
Y0 = torch.ones(n_v_full).to(gamma.device).float() * a0 * c0 |
|
norm = norm.view(-1, 3) |
|
nx, ny, nz = norm[:, 0], norm[:, 1], norm[:, 2] |
|
arrH = [] |
|
|
|
arrH.append(Y0) |
|
arrH.append(-a1 * c1 * ny) |
|
arrH.append(a1 * c1 * nz) |
|
arrH.append(-a1 * c1 * nx) |
|
arrH.append(a2 * c2 * nx * ny) |
|
arrH.append(-a2 * c2 * ny * nz) |
|
arrH.append(a2 * c2 * d0 * (3 * nz.pow(2) - 1)) |
|
arrH.append(-a2 * c2 * nx * nz) |
|
arrH.append(a2 * c2 * 0.5 * (nx.pow(2) - ny.pow(2))) |
|
|
|
H = torch.stack(arrH, 1) |
|
Y = H.view(n_b, num_vertex, 9) |
|
lighting = Y.bmm(gamma) |
|
|
|
face_color = face_texture * lighting |
|
return face_color |
|
|
|
def forward(self, rott_geometry, texture, diffuse_sh): |
|
face_normal = self.compute_normal(rott_geometry) |
|
face_color = self.Illumination_layer(texture, face_normal, diffuse_sh) |
|
face_color = TexturesVertex(face_color) |
|
mesh = Meshes( |
|
rott_geometry, |
|
self.tris.float().repeat(rott_geometry.shape[0], 1, 1), |
|
face_color, |
|
) |
|
rendered_img = self.renderer(mesh) |
|
rendered_img = torch.clamp(rendered_img, 0, 255) |
|
|
|
return rendered_img |
|
|
