import gc import torch import numpy as np from tqdm import tqdm import utils3d from PIL import Image from ..renderers import MeshRenderer, VoxelRenderer, PbrMeshRenderer from ..representations import Mesh, Voxel, MeshWithPbrMaterial, MeshWithVoxel from .random_utils import sphere_hammersley_sequence def yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs): is_list = isinstance(yaws, list) if not is_list: yaws = [yaws] pitchs = [pitchs] if not isinstance(rs, list): rs = [rs] * len(yaws) if not isinstance(fovs, list): fovs = [fovs] * len(yaws) extrinsics = [] intrinsics = [] for yaw, pitch, r, fov in zip(yaws, pitchs, rs, fovs): fov = torch.deg2rad(torch.tensor(float(fov))).cuda() yaw = torch.tensor(float(yaw)).cuda() pitch = torch.tensor(float(pitch)).cuda() orig = torch.tensor([ torch.sin(yaw) * torch.cos(pitch), torch.cos(yaw) * torch.cos(pitch), torch.sin(pitch), ]).cuda() * r extr = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda()) intr = utils3d.torch.intrinsics_from_fov_xy(fov, fov) extrinsics.append(extr) intrinsics.append(intr) if not is_list: extrinsics = extrinsics[0] intrinsics = intrinsics[0] return extrinsics, intrinsics def get_renderer(sample, **kwargs): if isinstance(sample, (MeshWithPbrMaterial, MeshWithVoxel)): renderer = PbrMeshRenderer() renderer.rendering_options.resolution = kwargs.get('resolution', 512) renderer.rendering_options.near = kwargs.get('near', 1) renderer.rendering_options.far = kwargs.get('far', 100) renderer.rendering_options.ssaa = kwargs.get('ssaa', 2) renderer.rendering_options.peel_layers = kwargs.get('peel_layers', 8) elif isinstance(sample, Mesh): renderer = MeshRenderer() renderer.rendering_options.resolution = kwargs.get('resolution', 512) renderer.rendering_options.near = kwargs.get('near', 1) renderer.rendering_options.far = kwargs.get('far', 100) renderer.rendering_options.ssaa = kwargs.get('ssaa', 2) renderer.rendering_options.chunk_size = kwargs.get('chunk_size', None) elif isinstance(sample, Voxel): renderer = VoxelRenderer() renderer.rendering_options.resolution = kwargs.get('resolution', 512) renderer.rendering_options.near = kwargs.get('near', 0.1) renderer.rendering_options.far = kwargs.get('far', 10.0) renderer.rendering_options.ssaa = kwargs.get('ssaa', 2) else: raise ValueError(f'Unsupported sample type: {type(sample)}') return renderer @torch.no_grad() def render_frames(sample, extrinsics, intrinsics, options={}, verbose=True, renderer=None, **kwargs): if renderer is None: renderer = get_renderer(sample, **options) rets = {} for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), total=len(extrinsics), desc='Rendering', disable=not verbose): res = renderer.render(sample, extr, intr, **kwargs) for k, v in res.items(): if k not in rets: rets[k] = [] if v.dim() == 2: v = v[None].repeat(3, 1, 1) rets[k].append(np.clip(v.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8)) return rets def render_video(sample, resolution=1024, bg_color=(0, 0, 0), num_frames=40, r=2, fov=40, start_yaw=None, start_pitch=None, **kwargs): """ Render a turntable video of the sample. Args: start_yaw: Starting yaw angle in radians. If None, defaults to π/2. start_pitch: Starting pitch angle in radians. If None, uses the default oscillating pitch starting at ~0.25. """ if start_yaw is None: start_yaw = np.pi / 2 yaws = -torch.linspace(0, 2 * 3.1415, num_frames) + start_yaw if start_pitch is not None: pitch = start_pitch + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames)) else: pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames)) yaws = yaws.tolist() pitch = pitch.tolist() extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov) return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) def render_multiview(sample, resolution=512, nviews=30): r = 2 fov = 40 cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)] yaws = [cam[0] for cam in cams] pitchs = [cam[1] for cam in cams] extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, r, fov) res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': (0, 0, 0)}) return res['color'], extrinsics, intrinsics def render_snapshot(samples, resolution=512, bg_color=(0, 0, 0), offset=(-16 / 180 * np.pi, 20 / 180 * np.pi), r=10, fov=8, nviews=4, **kwargs): yaw = np.linspace(0, 2 * np.pi, nviews, endpoint=False) yaw_offset = offset[0] yaw = [y + yaw_offset for y in yaw] pitch = [offset[1] for _ in range(nviews)] extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, r, fov) return render_frames(samples, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) def proj_camera_to_render_params(camera_angle_x, distance): """ Convert proj camera parameters to renderer-compatible extrinsics + intrinsics. The proj camera (Blender convention) views from the front. Through empirical testing, this corresponds to extrinsics_look_at from (0, 0, +distance) with up=Y in the mesh coordinate system used by the renderer. Args: camera_angle_x: horizontal FOV in radians distance: camera distance Returns: extrinsics: [4, 4] OpenCV world-to-camera (on CUDA) intrinsics: [3, 3] OpenCV normalized intrinsics (on CUDA) """ orig = torch.tensor([0.0, 0.0, distance]).cuda() target = torch.tensor([0.0, 0.0, 0.0]).cuda() up = torch.tensor([0.0, 1.0, 0.0]).cuda() extrinsics = utils3d.torch.extrinsics_look_at(orig, target, up) fov_tensor = torch.tensor(camera_angle_x, dtype=torch.float32).cuda() intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov_tensor, fov_tensor) return extrinsics, intrinsics def render_proj_aligned_video(sample, camera_angle_x, distance, resolution=1024, num_frames=40, bg_color=(0, 0, 0), **kwargs): """ Render a turntable video starting from the proj camera viewpoint. The first frame matches the proj input image exactly. Subsequent frames rotate around the object (around Y axis, which is up in mesh space). Args: sample: mesh to render camera_angle_x: proj camera FOV in radians distance: proj camera distance resolution: render resolution num_frames: number of video frames bg_color: background color **kwargs: additional kwargs (e.g. envmap) Returns: render result dict (same as render_frames) """ import math extr_first, intr_first = proj_camera_to_render_params(camera_angle_x, distance) extrinsics_list = [] intrinsics_list = [] angles = torch.linspace(0, 2 * math.pi, num_frames + 1)[:num_frames] for angle in angles: # Rotation around Y axis (up in mesh space) c = torch.cos(angle) s = torch.sin(angle) R_y = torch.tensor([ [ c, 0, s, 0], [ 0, 1, 0, 0], [-s, 0, c, 0], [ 0, 0, 0, 1], ], dtype=torch.float32).cuda() # world-to-camera for rotated world: extr @ R_y^{-1} R_y_inv = R_y.clone() R_y_inv[:3, :3] = R_y[:3, :3].T extr_rotated = extr_first @ R_y_inv extrinsics_list.append(extr_rotated) intrinsics_list.append(intr_first) return render_frames(sample, extrinsics_list, intrinsics_list, {'resolution': resolution, 'bg_color': bg_color}, **kwargs) def make_pbr_vis_frames(result, resolution=1024): num_frames = len(result['shaded']) frames = [] for i in range(num_frames): shaded = Image.fromarray(result['shaded'][i]) normal = Image.fromarray(result['normal'][i]) base_color = Image.fromarray(result['base_color'][i]) metallic = Image.fromarray(result['metallic'][i]) roughness = Image.fromarray(result['roughness'][i]) alpha = Image.fromarray(result['alpha'][i]) shaded = shaded.resize((resolution, resolution)) normal = normal.resize((resolution, resolution)) base_color = base_color.resize((resolution//2, resolution//2)) metallic = metallic.resize((resolution//2, resolution//2)) roughness = roughness.resize((resolution//2, resolution//2)) alpha = alpha.resize((resolution//2, resolution//2)) row1 = np.concatenate([shaded, normal], axis=1) row2 = np.concatenate([base_color, metallic, roughness, alpha], axis=1) frame = np.concatenate([row1, row2], axis=0) frames.append(frame) return frames