Spaces:
Running on Zero
Running on Zero
File size: 9,790 Bytes
a00e29b a6a6aa0 a00e29b 09e4a17 a00e29b 09e4a17 a00e29b 09e4a17 a00e29b a6a6aa0 a00e29b 09e4a17 a6a6aa0 | 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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | import torch
import numpy as np
from tqdm import tqdm
import utils3d
from easydict import EasyDict as edict
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_multiview_gpu(sample, resolution=512, nviews=10):
"""render_multiview on a black background, frames kept on the GPU as uint8 [N, H, W, 3]."""
cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)]
extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics([c[0] for c in cams], [c[1] for c in cams], 2, 40)
renderer = get_renderer(sample, resolution=resolution, bg_color=(0, 0, 0))
with torch.no_grad():
frames = torch.stack([
(renderer.render(sample, extr, intr)['color'].permute(1, 2, 0) * 255).clamp(0, 255).to(torch.uint8)
for extr, intr in zip(extrinsics, intrinsics)
])
return frames, 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)
class _FrozenGaussian:
def __init__(self, gaussian):
self.get_xyz = gaussian.get_xyz
self.get_opacity = gaussian.get_opacity
self.get_scaling = gaussian.get_scaling
self.get_rotation = gaussian.get_rotation
self.get_features = gaussian.get_features
self.active_sh_degree = gaussian.active_sh_degree
def _turntable_cameras(num_frames, resolution, r, fov, near, far):
from ..renderers.gaussian_render import intrinsics_to_projection
yaws = torch.linspace(0, 2 * 3.1415, num_frames)
pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws.tolist(), pitch.tolist(), r, fov)
views = torch.stack(extrinsics)
centers = torch.inverse(views)[:, :3, 3]
fovs = (2 * torch.atan(0.5 / torch.stack([intr[0, 0] for intr in intrinsics]))).tolist()
cameras = []
for view, intr, center, fov_rad in zip(views, intrinsics, centers, fovs):
perspective = intrinsics_to_projection(intr, near, far)
cameras.append(edict({
"image_height": resolution, "image_width": resolution, "FoVx": fov_rad, "FoVy": fov_rad,
"znear": near, "zfar": far, "world_view_transform": view.T.contiguous(),
"projection_matrix": perspective.T.contiguous(),
"full_proj_transform": (perspective @ view).T.contiguous(), "camera_center": center,
}))
return cameras
def render_gaussian_turntable(sample, resolution, num_frames, bg_colors, r=2, fov=40):
"""Same frames as render_video(mode='color') per background, kept on the GPU as uint8 [N, H, W, 3]."""
from ..renderers.gaussian_render import render
with torch.no_grad():
renderer = get_renderer(sample, resolution=resolution)
options = renderer.rendering_options
cameras = _turntable_cameras(num_frames, resolution, r, fov, options.near, options.far)
frozen = _FrozenGaussian(sample)
backgrounds = [torch.tensor(bg, dtype=torch.float32, device="cuda") for bg in bg_colors]
frames = torch.empty(len(bg_colors), num_frames, resolution, resolution, 3, dtype=torch.uint8, device="cuda")
for j, camera in enumerate(cameras):
for b, background in enumerate(backgrounds):
color = render(camera, frozen, renderer.pipe, background, scaling_modifier=renderer.pipe.scale_modifier).render
frames[b, j] = (color.permute(1, 2, 0) * 255).clamp(0, 255).to(torch.uint8)
return frames
|