"""
Author: Yao Feng
Copyright (c) 2020, Yao Feng
All rights reserved.
"""
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from skimage.io import imread
import imageio
from . import util
def set_rasterizer(type='pytorch3d'):
if type == 'pytorch3d':
global Meshes, load_obj, rasterize_meshes
from pytorch3d.structures import Meshes
from pytorch3d.io import load_obj
from pytorch3d.renderer.mesh import rasterize_meshes
elif type == 'standard':
global standard_rasterize, load_obj
import os
from .util import load_obj
# Use JIT Compiling Extensions
# ref: https://pytorch.org/tutorials/advanced/cpp_extension.html
from torch.utils.cpp_extension import load, CUDA_HOME
curr_dir = os.path.dirname(__file__)
standard_rasterize_cuda = \
load(name='standard_rasterize_cuda',
sources=[f'{curr_dir}/rasterizer/standard_rasterize_cuda.cpp',
f'{curr_dir}/rasterizer/standard_rasterize_cuda_kernel.cu'],
extra_cuda_cflags=['-std=c++14', '-ccbin=$$(which gcc-7)']) # cuda10.2 is not compatible with gcc9. Specify gcc 7
from standard_rasterize_cuda import standard_rasterize
# If JIT does not work, try manually installation first
# 1. see instruction here: pixielib/utils/rasterizer/INSTALL.md
# 2. add this: "from .rasterizer.standard_rasterize_cuda import standard_rasterize" here
class StandardRasterizer(nn.Module):
""" Alg: https://www.scratchapixel.com/lessons/3d-basic-rendering/rasterization-practical-implementation
Notice:
x,y,z are in image space, normalized to [-1, 1]
can render non-squared image
not differentiable
"""
def __init__(self, height, width=None):
"""
use fixed raster_settings for rendering faces
"""
super().__init__()
if width is None:
width = height
self.h = h = height
self.w = w = width
def forward(self, vertices, faces, attributes=None, h=None, w=None):
device = vertices.device
if h is None:
h = self.h
if w is None:
w = self.h
bz = vertices.shape[0]
depth_buffer = torch.zeros([bz, h, w]).float().to(device) + 1e6
triangle_buffer = torch.zeros([bz, h, w]).int().to(device) - 1
baryw_buffer = torch.zeros([bz, h, w, 3]).float().to(device)
vert_vis = torch.zeros([bz, vertices.shape[1]]).float().to(device)
vertices = vertices.clone().float()
vertices[..., 0] = vertices[..., 0] * w / 2 + w / 2
vertices[..., 1] = vertices[..., 1] * h / 2 + h / 2
vertices[..., 2] = vertices[..., 2] * w / 2
f_vs = util.face_vertices(vertices, faces)
standard_rasterize(f_vs, depth_buffer, triangle_buffer, baryw_buffer,
h, w)
pix_to_face = triangle_buffer[:, :, :, None].long()
bary_coords = baryw_buffer[:, :, :, None, :]
vismask = (pix_to_face > -1).float()
D = attributes.shape[-1]
attributes = attributes.clone()
attributes = attributes.view(attributes.shape[0] * attributes.shape[1],
3, attributes.shape[-1])
N, H, W, K, _ = bary_coords.shape
mask = pix_to_face == -1
pix_to_face = pix_to_face.clone()
pix_to_face[mask] = 0
idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D)
pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D)
pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2)
pixel_vals[mask] = 0 # Replace masked values in output.
pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2)
pixel_vals = torch.cat(
[pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1)
return pixel_vals
class Pytorch3dRasterizer(nn.Module):
""" Borrowed from https://github.com/facebookresearch/pytorch3d
This class implements methods for rasterizing a batch of heterogenous Meshes.
Notice:
x,y,z are in image space, normalized
can only render squared image now
"""
def __init__(self, image_size=224):
"""
use fixed raster_settings for rendering faces
"""
super().__init__()
raster_settings = {
'image_size': image_size,
'blur_radius': 0.0,
'faces_per_pixel': 1,
'bin_size': None,
'max_faces_per_bin': None,
'perspective_correct': False,
}
raster_settings = util.dict2obj(raster_settings)
self.raster_settings = raster_settings
def forward(self, vertices, faces, attributes=None, h=None, w=None):
fixed_vertices = vertices.clone()
fixed_vertices[..., :2] = -fixed_vertices[..., :2]
meshes_screen = Meshes(verts=fixed_vertices.float(),
faces=faces.long())
raster_settings = self.raster_settings
pix_to_face, zbuf, bary_coords, dists = rasterize_meshes(
meshes_screen,
image_size=raster_settings.image_size,
blur_radius=raster_settings.blur_radius,
faces_per_pixel=raster_settings.faces_per_pixel,
bin_size=raster_settings.bin_size,
max_faces_per_bin=raster_settings.max_faces_per_bin,
perspective_correct=raster_settings.perspective_correct,
)
vismask = (pix_to_face > -1).float()
D = attributes.shape[-1]
attributes = attributes.clone()
attributes = attributes.view(attributes.shape[0] * attributes.shape[1],
3, attributes.shape[-1])
N, H, W, K, _ = bary_coords.shape
mask = pix_to_face == -1
pix_to_face = pix_to_face.clone()
pix_to_face[mask] = 0
idx = pix_to_face.view(N * H * W * K, 1, 1).expand(N * H * W * K, 3, D)
pixel_face_vals = attributes.gather(0, idx).view(N, H, W, K, 3, D)
pixel_vals = (bary_coords[..., None] * pixel_face_vals).sum(dim=-2)
pixel_vals[mask] = 0 # Replace masked values in output.
pixel_vals = pixel_vals[:, :, :, 0].permute(0, 3, 1, 2)
pixel_vals = torch.cat(
[pixel_vals, vismask[:, :, :, 0][:, None, :, :]], dim=1)
return pixel_vals
class SRenderY(nn.Module):
def __init__(self,
image_size,
obj_filename,
uv_size=256,
rasterizer_type='standard'):
super(SRenderY, self).__init__()
self.image_size = image_size
self.uv_size = uv_size
if rasterizer_type == 'pytorch3d':
self.rasterizer = Pytorch3dRasterizer(image_size)
self.uv_rasterizer = Pytorch3dRasterizer(uv_size)
verts, faces, aux = load_obj(obj_filename)
uvcoords = aux.verts_uvs[None, ...] # (N, V, 2)
uvfaces = faces.textures_idx[None, ...] # (N, F, 3)
faces = faces.verts_idx[None, ...]
elif rasterizer_type == 'standard':
self.rasterizer = StandardRasterizer(image_size)
self.uv_rasterizer = StandardRasterizer(uv_size)
verts, uvcoords, faces, uvfaces = load_obj(obj_filename)
verts = verts[None, ...]
uvcoords = uvcoords[None, ...]
faces = faces[None, ...]
uvfaces = uvfaces[None, ...]
else:
NotImplementedError
# faces
dense_triangles = util.generate_triangles(uv_size, uv_size)
self.register_buffer(
'dense_faces',
torch.from_numpy(dense_triangles).long()[None, :, :])
self.register_buffer('faces', faces)
self.register_buffer('raw_uvcoords', uvcoords)
# uv coords
uvcoords = torch.cat([uvcoords, uvcoords[:, :, 0:1] * 0. + 1.],
-1) # [bz, ntv, 3]
uvcoords = uvcoords * 2 - 1
uvcoords[..., 1] = -uvcoords[..., 1]
face_uvcoords = util.face_vertices(uvcoords, uvfaces)
self.register_buffer('uvcoords', uvcoords)
self.register_buffer('uvfaces', uvfaces)
self.register_buffer('face_uvcoords', face_uvcoords)
# shape colors, for rendering shape overlay
colors = torch.tensor([180, 180, 180])[None, None, :].repeat(
1,
faces.max() + 1, 1).float() / 255.
face_colors = util.face_vertices(colors, faces)
self.register_buffer('vertex_colors', colors)
self.register_buffer('face_colors', face_colors)
# SH factors for lighting
pi = np.pi
constant_factor = torch.tensor([
1 / np.sqrt(4 * pi), ((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))),
((2 * pi) / 3) * (np.sqrt(3 / (4 * pi))), ((2 * pi) / 3) *
(np.sqrt(3 / (4 * pi))), (pi / 4) * (3) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (3) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (3) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (3 / 2) * (np.sqrt(5 / (12 * pi))),
(pi / 4) * (1 / 2) * (np.sqrt(5 / (4 * pi)))
]).float()
self.register_buffer('constant_factor', constant_factor)
def forward(self,
vertices,
transformed_vertices,
albedos,
lights=None,
light_type='point',
background=None,
h=None,
w=None):
'''
-- Texture Rendering
vertices: [batch_size, V, 3], vertices in world space, for calculating normals, then shading
transformed_vertices: [batch_size, V, 3], rnage:[-1,1], projected vertices, in image space, for rasterization
albedos: [batch_size, 3, h, w], uv map
lights:
spherical homarnic: [N, 9(shcoeff), 3(rgb)]
points/directional lighting: [N, n_lights, 6(xyzrgb)]
light_type:
point or directional
'''
batch_size = vertices.shape[0]
# normalize z to 10-90 for raterization (in pytorch3d, near far: 0-100)
transformed_vertices = transformed_vertices.clone()
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] - transformed_vertices[:, :,
2].min(
)
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] / transformed_vertices[:, :,
2].max(
)
transformed_vertices[:, :, 2] = transformed_vertices[:, :, 2] * 80 + 10
# attributes
face_vertices = util.face_vertices(
vertices, self.faces.expand(batch_size, -1, -1))
normals = util.vertex_normals(vertices,
self.faces.expand(batch_size, -1, -1))
face_normals = util.face_vertices(
normals, self.faces.expand(batch_size, -1, -1))
transformed_normals = util.vertex_normals(
transformed_vertices, self.faces.expand(batch_size, -1, -1))
transformed_face_normals = util.face_vertices(
transformed_normals, self.faces.expand(batch_size, -1, -1))
attributes = torch.cat([
self.face_uvcoords.expand(batch_size, -1, -1, -1),
transformed_face_normals.detach(),
face_vertices.detach(), face_normals
], -1)
# rasterize
rendering = self.rasterizer(transformed_vertices,
self.faces.expand(batch_size, -1, -1),
attributes, h, w)
####
# vis mask
alpha_images = rendering[:, -1, :, :][:, None, :, :].detach()
# albedo
uvcoords_images = rendering[:, :3, :, :]
grid = (uvcoords_images).permute(0, 2, 3, 1)[:, :, :, :2]
albedo_images = F.grid_sample(albedos, grid, align_corners=False)
# visible mask for pixels with positive normal direction
transformed_normal_map = rendering[:, 3:6, :, :].detach()
pos_mask = (transformed_normal_map[:, 2:, :, :] < -0.05).float()
# shading
normal_images = rendering[:, 9:12, :, :]
if lights is not None:
if lights.shape[1] == 9:
shading_images = self.add_SHlight(normal_images, lights)
else:
if light_type == 'point':
vertice_images = rendering[:, 6:9, :, :].detach()
shading = self.add_pointlight(
vertice_images.permute(0, 2, 3,
1).reshape([batch_size, -1, 3]),
normal_images.permute(0, 2, 3,
1).reshape([batch_size, -1, 3]),
lights)
shading_images = shading.reshape([
batch_size, albedo_images.shape[2],
albedo_images.shape[3], 3
]).permute(0, 3, 1, 2)
else:
shading = self.add_directionlight(
normal_images.permute(0, 2, 3,
1).reshape([batch_size, -1, 3]),
lights)
shading_images = shading.reshape([
batch_size, albedo_images.shape[2],
albedo_images.shape[3], 3
]).permute(0, 3, 1, 2)
images = albedo_images * shading_images
else:
images = albedo_images
shading_images = images.detach() * 0.
if background is None:
images = images*alpha_images + \
torch.ones_like(images).to(vertices.device)*(1-alpha_images)
else:
# background = F.interpolate(background, [self.image_size, self.image_size])
images = images * alpha_images + background.contiguous() * (
1 - alpha_images)
outputs = {
'images': images,
'albedo_images': albedo_images,
'alpha_images': alpha_images,
'pos_mask': pos_mask,
'shading_images': shading_images,
'grid': grid,
'normals': normals,
'normal_images': normal_images,
'transformed_normals': transformed_normals,
}
return outputs
def add_SHlight(self, normal_images, sh_coeff):
'''
sh_coeff: [bz, 9, 3]
'''
N = normal_images
sh = torch.stack([
N[:, 0] * 0. + 1., N[:, 0], N[:, 1], N[:, 2], N[:, 0] * N[:, 1],
N[:, 0] * N[:, 2], N[:, 1] * N[:, 2], N[:, 0]**2 - N[:, 1]**2, 3 *
(N[:, 2]**2) - 1
], 1) # [bz, 9, h, w]
sh = sh * self.constant_factor[None, :, None, None]
# [bz, 9, 3, h, w]
shading = torch.sum(
sh_coeff[:, :, :, None, None] * sh[:, :, None, :, :], 1)
return shading
def add_pointlight(self, vertices, normals, lights):
'''
vertices: [bz, nv, 3]
lights: [bz, nlight, 6]
returns:
shading: [bz, nv, 3]
'''
light_positions = lights[:, :, :3]
light_intensities = lights[:, :, 3:]
directions_to_lights = F.normalize(light_positions[:, :, None, :] -
vertices[:, None, :, :],
dim=3)
# normals_dot_lights = torch.clamp((normals[:,None,:,:]*directions_to_lights).sum(dim=3), 0., 1.)
normals_dot_lights = (normals[:, None, :, :] *
directions_to_lights).sum(dim=3)
shading = normals_dot_lights[:, :, :,
None] * light_intensities[:, :, None, :]
return shading.mean(1)
def add_directionlight(self, normals, lights):
'''
normals: [bz, nv, 3]
lights: [bz, nlight, 6]
returns:
shading: [bz, nv, 3]
'''
light_direction = lights[:, :, :3]
light_intensities = lights[:, :, 3:]
directions_to_lights = F.normalize(
light_direction[:, :, None, :].expand(-1, -1, normals.shape[1],
-1),
dim=3)
# normals_dot_lights = torch.clamp((normals[:,None,:,:]*directions_to_lights).sum(dim=3), 0., 1.)
# normals_dot_lights = (normals[:,None,:,:]*directions_to_lights).sum(dim=3)
normals_dot_lights = torch.clamp(
(normals[:, None, :, :] * directions_to_lights).sum(dim=3), 0., 1.)
shading = normals_dot_lights[:, :, :,
None] * light_intensities[:, :, None, :]
return shading.mean(1)
def render_shape(self,
vertices,
transformed_vertices,
colors=None,
background=None,
detail_normal_images=None,
lights=None,
return_grid=False,
uv_detail_normals=None,
h=None,
w=None):
'''
-- rendering shape with detail normal map
'''
batch_size = vertices.shape[0]
if lights is None:
light_positions = torch.tensor([
[-5, 5, -5],
[5, 5, -5],
[-5, -5, -5],
[5, -5, -5],
[0, 0, -5],
])[None, :, :].expand(batch_size, -1, -1).float()
light_intensities = torch.ones_like(light_positions).float() * 1.7
lights = torch.cat((light_positions, light_intensities),
2).to(vertices.device)
# normalize z to 10-90 for raterization (in pytorch3d, near far: 0-100)
transformed_vertices = transformed_vertices.clone()
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] - transformed_vertices[:, :,
2].min(
)
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] / transformed_vertices[:, :,
2].max(
)
transformed_vertices[:, :, 2] = transformed_vertices[:, :, 2] * 80 + 10
# Attributes
face_vertices = util.face_vertices(
vertices, self.faces.expand(batch_size, -1, -1))
normals = util.vertex_normals(vertices,
self.faces.expand(batch_size, -1, -1))
face_normals = util.face_vertices(
normals, self.faces.expand(batch_size, -1, -1))
transformed_normals = util.vertex_normals(
transformed_vertices, self.faces.expand(batch_size, -1, -1))
transformed_face_normals = util.face_vertices(
transformed_normals, self.faces.expand(batch_size, -1, -1))
if colors is None:
colors = self.face_colors.expand(batch_size, -1, -1, -1)
attributes = torch.cat([
colors,
transformed_face_normals.detach(),
face_vertices.detach(), face_normals,
self.face_uvcoords.expand(batch_size, -1, -1, -1)
], -1)
# rasterize
rendering = self.rasterizer(transformed_vertices,
self.faces.expand(batch_size, -1, -1),
attributes, h, w)
####
alpha_images = rendering[:, -1, :, :][:, None, :, :].detach()
# albedo
albedo_images = rendering[:, :3, :, :]
# mask
transformed_normal_map = rendering[:, 3:6, :, :].detach()
pos_mask = (transformed_normal_map[:, 2:, :, :] < 0).float()
# shading
normal_images = rendering[:, 9:12, :, :].detach()
vertice_images = rendering[:, 6:9, :, :].detach()
if detail_normal_images is not None:
normal_images = detail_normal_images
if uv_detail_normals is not None:
uvcoords_images = rendering[:, 12:15, :, :]
grid = (uvcoords_images).permute(0, 2, 3, 1)[:, :, :, :2]
detail_normal_images = F.grid_sample(uv_detail_normals,
grid,
align_corners=False)
normal_images = detail_normal_images
shading = self.add_directionlight(
normal_images.permute(0, 2, 3, 1).reshape([batch_size, -1, 3]),
lights)
shading_images = shading.reshape(
[batch_size, albedo_images.shape[2], albedo_images.shape[3],
3]).permute(0, 3, 1, 2).contiguous()
shaded_images = albedo_images * shading_images
if background is None:
shape_images = shaded_images*alpha_images + \
torch.ones_like(shaded_images).to(
vertices.device)*(1-alpha_images)
else:
# background = F.interpolate(background, [self.image_size, self.image_size])
shape_images = shaded_images*alpha_images + \
background.contiguous()*(1-alpha_images)
if return_grid:
uvcoords_images = rendering[:, 12:15, :, :]
grid = (uvcoords_images).permute(0, 2, 3, 1)[:, :, :, :2]
return shape_images, normal_images, grid
else:
return shape_images
def render_depth(self, transformed_vertices):
'''
-- rendering depth
'''
transformed_vertices = transformed_vertices.clone()
batch_size = transformed_vertices.shape[0]
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] - transformed_vertices[:, :,
2].min(
)
z = -transformed_vertices[:, :, 2:].repeat(1, 1, 3)
z = z - z.min()
z = z / z.max()
# Attributes
attributes = util.face_vertices(z,
self.faces.expand(batch_size, -1, -1))
# rasterize
rendering = self.rasterizer(transformed_vertices,
self.faces.expand(batch_size, -1, -1),
attributes)
####
alpha_images = rendering[:, -1, :, :][:, None, :, :].detach()
depth_images = rendering[:, :1, :, :]
return depth_images
def render_colors(self, transformed_vertices, colors, h=None, w=None):
'''
-- rendering colors: could be rgb color/ normals, etc
colors: [bz, num of vertices, 3]
'''
transformed_vertices = transformed_vertices.clone()
batch_size = colors.shape[0]
# normalize z to 10-90 for raterization (in pytorch3d, near far: 0-100)
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] - transformed_vertices[:, :,
2].min(
)
transformed_vertices[:, :,
2] = transformed_vertices[:, :,
2] / transformed_vertices[:, :,
2].max(
)
transformed_vertices[:, :, 2] = transformed_vertices[:, :, 2] * 80 + 10
# Attributes
attributes = util.face_vertices(colors,
self.faces.expand(batch_size, -1, -1))
# rasterize
rendering = self.rasterizer(transformed_vertices,
self.faces.expand(batch_size, -1, -1),
attributes,
h=h,
w=w)
####
alpha_images = rendering[:, [-1], :, :].detach()
images = rendering[:, :3, :, :] * alpha_images
return images
def world2uv(self, vertices):
'''
project vertices from world space to uv space
vertices: [bz, V, 3]
uv_vertices: [bz, 3, h, w]
'''
batch_size = vertices.shape[0]
face_vertices = util.face_vertices(
vertices, self.faces.expand(batch_size, -1, -1))
uv_vertices = self.uv_rasterizer(
self.uvcoords.expand(batch_size, -1, -1),
self.uvfaces.expand(batch_size, -1, -1), face_vertices)[:, :3]
return uv_vertices