Spaces:
Running on Zero
Running on Zero
File size: 6,491 Bytes
a00e29b 09e4a17 a00e29b 09e4a17 a00e29b 09e4a17 a00e29b 09e4a17 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | import torch
import numpy as np
from tqdm import tqdm
import utils3d
from ..renderers import GaussianRenderer, MeshRenderer
from ..representations import Gaussian, MeshExtractResult
# from ..modules import sparse as sp
from .random_utils import sphere_hammersley_sequence
def yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs):
with torch.no_grad():
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, Gaussian):
renderer = GaussianRenderer()
renderer.rendering_options.resolution = kwargs.get('resolution', 512)
renderer.rendering_options.near = kwargs.get('near', 0.8)
renderer.rendering_options.far = kwargs.get('far', 1.6)
renderer.rendering_options.bg_color = kwargs.get('bg_color', (0, 0, 0))
renderer.rendering_options.ssaa = kwargs.get('ssaa', 1)
renderer.pipe.kernel_size = kwargs.get('kernel_size', 0.1)
renderer.pipe.use_mip_gaussian = True
elif isinstance(sample, MeshExtractResult):
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', 1)
else:
raise ValueError(f'Unsupported sample type: {type(sample)}')
return renderer
def render_frames(sample, extrinsics, intrinsics, options={}, colors_overwrite=None, verbose=True, **kwargs):
renderer = get_renderer(sample, **options)
rets = {}
mode = kwargs.get('mode')
for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), desc='Rendering', disable=not verbose):
if isinstance(sample, MeshExtractResult):
has_vertex_color = getattr(sample, "vertex_attrs", None) is not None
return_types = ["normal"]
if mode == "color":
return_types = ["color"] if has_vertex_color else ["normal"]
elif mode == "normal":
return_types = ["normal"]
else:
# Default behavior for mesh keeps both channels available.
return_types = ["color", "normal"] if has_vertex_color else ["normal"]
res = renderer.render(sample, extr, intr, return_types=return_types)
if 'color' in res:
if 'color' not in rets:
rets['color'] = []
rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
if 'normal' in res:
if 'normal' not in rets:
rets['normal'] = []
rets['normal'].append(np.clip(res['normal'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
if mode == "color" and 'color' not in rets and 'normal' in rets:
# Fallback to normal frames when no vertex colors exist.
rets['color'] = list(rets['normal'])
else:
res = renderer.render(sample, extr, intr, colors_overwrite=colors_overwrite)
if 'color' not in rets: rets['color'] = []
if 'depth' not in rets: rets['depth'] = []
rets['color'].append(np.clip(res['color'].detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
if 'percent_depth' in res:
rets['depth'].append(res['percent_depth'].detach().cpu().numpy())
elif 'depth' in res:
rets['depth'].append(res['depth'].detach().cpu().numpy())
else:
rets['depth'].append(None)
return rets
def render_video(sample, resolution=512, bg_color=(0, 0, 0), num_frames=300, r=2, fov=40, **kwargs):
with torch.no_grad():
yaws = torch.linspace(0, 2 * 3.1415, num_frames)
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=10):
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)})
if 'color' in res:
return res['color'], extrinsics, intrinsics
if 'normal' in res:
return res['normal'], extrinsics, intrinsics
return [], 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, **kwargs):
yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
yaw_offset = offset[0]
yaw = [y + yaw_offset for y in yaw]
pitch = [offset[1] for _ in range(4)]
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)
|