255 lines
8.1 KiB
Python
255 lines
8.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import numpy as np
|
|
|
|
from pytorch3d.renderer import (
|
|
RasterizationSettings,
|
|
MeshRenderer,
|
|
MeshRasterizer,
|
|
HardPhongShader,
|
|
PointLights,
|
|
PerspectiveCameras,
|
|
)
|
|
from pytorch3d.structures import Meshes
|
|
from pytorch3d.structures.meshes import join_meshes_as_scene
|
|
from pytorch3d.renderer import TexturesAtlas, TexturesVertex
|
|
from pytorch3d.transforms import euler_angles_to_matrix
|
|
|
|
from .tools import get_colors, read_image, transform_torch3d, checkerboard_geometry
|
|
|
|
|
|
def init_renderer(img_size, intrins, device, vis_scale=1.0, bg_paths=None):
|
|
img_size = int(vis_scale * img_size[0]), int(vis_scale * img_size[1])
|
|
renderer = RenderBase(device, img_size, intrins=vis_scale * intrins)
|
|
if bg_paths is not None:
|
|
bg_imgs = [
|
|
torch.from_numpy(read_image(p, scale=vis_scale)).float() / 255
|
|
for p in bg_paths
|
|
]
|
|
print(bg_imgs[0].shape)
|
|
renderer.set_bg_seq(bg_imgs)
|
|
print("RENDERER", renderer)
|
|
return renderer
|
|
|
|
|
|
class RenderBase(nn.Module):
|
|
def __init__(
|
|
self,
|
|
device,
|
|
img_size,
|
|
intrins=None,
|
|
blur_radius: float = 0.0, # rasterizer settings
|
|
faces_per_pixel: int = 3,
|
|
):
|
|
"""
|
|
Uses Pytorch3D to render results.
|
|
We use the PyTorch3D coordinate system, which assumes:
|
|
+X:left, +Y: up and +Z: from us to scene (right-handed)
|
|
https://pytorch3d.org/docs/cameras
|
|
"""
|
|
super().__init__()
|
|
self.device = device
|
|
W, H = img_size
|
|
self.img_size = (H, W)
|
|
if intrins is not None:
|
|
intrins = intrins.to(device)
|
|
self.cam_f, self.cam_center = intrins[:2], intrins[2:]
|
|
cx, cy = self.cam_center.detach().cpu()
|
|
else:
|
|
self.cam_f = torch.tensor([0.5 * (H + W), 0.5 * (H + W)], device=device)
|
|
self.cam_center = torch.tensor([0.5 * W, 0.5 * H], device=device)
|
|
print("cam_f", self.cam_f, "cam_center", self.cam_center)
|
|
print("img_size", self.img_size)
|
|
self.blur_radius = blur_radius
|
|
self.faces_per_pixel = faces_per_pixel
|
|
|
|
self.light_loc = [[0, 0, -0.5]]
|
|
self.cameras = self.make_cameras()
|
|
self.bg_imgs = None
|
|
self.set_ground()
|
|
self.build_renderer()
|
|
|
|
def build_renderer(self):
|
|
self.lights = PointLights(device=self.device, location=self.light_loc)
|
|
|
|
self.rasterizer = RasterizationSettings(
|
|
image_size=self.img_size,
|
|
blur_radius=self.blur_radius,
|
|
faces_per_pixel=self.faces_per_pixel,
|
|
max_faces_per_bin=100000,
|
|
)
|
|
self.renderer = MeshRenderer(
|
|
rasterizer=MeshRasterizer(
|
|
cameras=self.cameras, raster_settings=self.rasterizer
|
|
),
|
|
shader=HardPhongShader(
|
|
cameras=self.cameras, lights=self.lights, device=self.device
|
|
),
|
|
)
|
|
|
|
def make_cameras(self, cam_poses=None):
|
|
device = self.device
|
|
if cam_poses is None:
|
|
cam_poses = torch.eye(4, device=device)[None]
|
|
cam_poses = cam_poses.to(device)
|
|
cam_R, cam_t = transform_torch3d(cam_poses)
|
|
Nc = len(cam_poses)
|
|
focal = self.cam_f[None].repeat(Nc, 1)
|
|
ppt = self.cam_center[None].repeat(Nc, 1)
|
|
img_size = [self.img_size for _ in range(Nc)]
|
|
return PerspectiveCameras(
|
|
device=self.device,
|
|
in_ndc=False,
|
|
focal_length=focal,
|
|
principal_point=ppt,
|
|
image_size=img_size,
|
|
R=cam_R,
|
|
T=cam_t,
|
|
)
|
|
|
|
def render_frame(self, cam_poses, verts, faces, colors):
|
|
"""
|
|
:param cam_poses (C, 4, 4)
|
|
:param verts (B, V, 3)
|
|
:param faces (F, 3)
|
|
:param colors (B, 4)
|
|
"""
|
|
with torch.no_grad():
|
|
cameras = self.make_cameras(cam_poses)
|
|
mesh = build_meshes(verts, faces, colors)
|
|
image = self.renderer(mesh, cameras=cameras)
|
|
return image
|
|
|
|
def render_with_ground(self, cam_poses, verts, faces, colors, ground_pose):
|
|
"""
|
|
:param cam_poses (C, 4, 4)
|
|
:param verts (B, V, 3)
|
|
:param faces (F, 3)
|
|
:param colors (B, 3)
|
|
:param ground_pose (4, 4)
|
|
"""
|
|
with torch.no_grad():
|
|
self.set_ground_pose(ground_pose)
|
|
cameras = self.make_cameras(cam_poses)
|
|
# (B, V, 3), (B, F, 3), (B, V, 3)
|
|
verts, faces, colors = prep_shared_geometry(verts, faces, colors)
|
|
# (V, 3), (F, 3), (V, 3)
|
|
gv, gf, gc = self.ground_geometry
|
|
verts = list(torch.unbind(verts, dim=0)) + [gv]
|
|
faces = list(torch.unbind(faces, dim=0)) + [gf]
|
|
colors = list(torch.unbind(colors, dim=0)) + [gc[..., :3]]
|
|
mesh = create_meshes(verts, faces, colors)
|
|
|
|
image = self.renderer(mesh, cameras=cameras)
|
|
return image
|
|
|
|
def set_bg_seq(self, bg_imgs):
|
|
self.bg_imgs = bg_imgs
|
|
|
|
def set_ground(self, pose=None):
|
|
device = self.device
|
|
v, f, vc, fc = map(torch.from_numpy, checkerboard_geometry(up="y"))
|
|
v, f, vc = v.to(device), f.to(device), vc.to(device)
|
|
self.ground_geometry = [v, f, vc]
|
|
if pose is not None:
|
|
self.set_ground_pose(pose)
|
|
|
|
def set_ground_pose(self, pose):
|
|
"""
|
|
:param pose (4, 4)
|
|
"""
|
|
if self.ground_geometry is None:
|
|
return
|
|
v = self.ground_geometry[0]
|
|
pose = pose.to(self.device)
|
|
# R, t = pose[:3, :3], pose[:3, 3]
|
|
R, t = transform_torch3d(pose)
|
|
self.ground_geometry[0] = torch.einsum("ij,vj->vi", R, v) + t[None]
|
|
|
|
def composite_bg(self, frames, fac=1.0):
|
|
T = len(frames)
|
|
if self.bg_imgs is None:
|
|
return frames
|
|
|
|
out_frames = []
|
|
for bg, frame in zip(self.bg_imgs, frames):
|
|
alpha = fac * frame[..., 3:]
|
|
out = alpha * frame[..., :3] + (1 - alpha) * bg
|
|
out_frames.append(out)
|
|
return out_frames
|
|
|
|
def render_video(
|
|
self,
|
|
cam_poses,
|
|
verts,
|
|
faces,
|
|
colors,
|
|
render_bg=False,
|
|
fac=1.0,
|
|
ground_pose=None,
|
|
):
|
|
"""
|
|
:param cam_poses (C, T, 4, 4)
|
|
:param verts list of T (B, V, 3)
|
|
:param faces list of T (F, 3)
|
|
:param colors list of T (B, 4)
|
|
:param render_bg (optional default False)
|
|
:param fac (optional float, default 1)
|
|
:param ground_pose (optional, default None) (4, 4)
|
|
returns list of T (H, W, 4)
|
|
"""
|
|
T = len(verts)
|
|
print(f"rendering video with {T} frames")
|
|
if cam_poses.shape[1] == 1:
|
|
cam_poses = cam_poses.expand(-1, T, -1, -1)
|
|
|
|
if render_bg: # don't render ground on source frame
|
|
ground_pose = None
|
|
|
|
frames = []
|
|
for t in range(T):
|
|
if ground_pose is None:
|
|
# (H, W, 4)
|
|
frame = self.render_frame(
|
|
cam_poses[:, t], verts[t], faces[t], colors[t]
|
|
)
|
|
else:
|
|
# (H, W, 4)
|
|
frame = self.render_with_ground(
|
|
cam_poses[:, t], verts[t], faces[t], colors[t], ground_pose
|
|
)
|
|
frames.append(frame.cpu()[0])
|
|
|
|
if render_bg and self.bg_imgs is not None:
|
|
frames = self.composite_bg(frames, fac=fac)
|
|
|
|
return [(255 * f).byte() for f in frames]
|
|
|
|
|
|
def build_meshes(verts, faces, colors):
|
|
return create_meshes(*prep_shared_geometry(verts, faces, colors))
|
|
|
|
|
|
def prep_shared_geometry(verts, faces, colors):
|
|
"""
|
|
:param verts (B, V, 3)
|
|
:param faces (F, 3)
|
|
:param colors (B, 4)
|
|
"""
|
|
B, V, _ = verts.shape
|
|
F, _ = faces.shape
|
|
colors = colors.unsqueeze(1).expand(B, V, -1)[..., :3]
|
|
faces = faces.unsqueeze(0).expand(B, F, -1)
|
|
return verts, faces, colors
|
|
|
|
|
|
def create_meshes(verts, faces, colors):
|
|
"""
|
|
:param verts (B, V, 3)
|
|
:param faces (B, F, 3)
|
|
:param colors (B, V, 3)
|
|
"""
|
|
textures = TexturesVertex(verts_features=colors)
|
|
meshes = Meshes(verts=verts, faces=faces, textures=textures)
|
|
return join_meshes_as_scene(meshes)
|