a fork of shap-e for gc
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

290 lines
10 KiB

2 years ago
from functools import partial
from typing import Any, Dict, Optional, Sequence, Tuple, Union
import torch
from shap_e.models.nerf.model import NeRFModel
from shap_e.models.nerf.ray import RayVolumeIntegral, StratifiedRaySampler, render_rays
from shap_e.models.nn.meta import subdict
from shap_e.models.nn.utils import to_torch
from shap_e.models.query import Query
from shap_e.models.renderer import RayRenderer, render_views_from_rays
from shap_e.models.stf.base import Model
from shap_e.models.stf.renderer import STFRendererBase, render_views_from_stf
from shap_e.models.volume import BoundingBoxVolume, Volume
from shap_e.rendering.blender.constants import BASIC_AMBIENT_COLOR, BASIC_DIFFUSE_COLOR
from shap_e.util.collections import AttrDict
class NeRSTFRenderer(RayRenderer, STFRendererBase):
def __init__(
self,
sdf: Optional[Model],
tf: Optional[Model],
nerstf: Optional[Model],
void: NeRFModel,
volume: Volume,
grid_size: int,
n_coarse_samples: int,
n_fine_samples: int,
importance_sampling_options: Optional[Dict[str, Any]] = None,
separate_shared_samples: bool = False,
texture_channels: Sequence[str] = ("R", "G", "B"),
channel_scale: Sequence[float] = (255.0, 255.0, 255.0),
ambient_color: Union[float, Tuple[float]] = BASIC_AMBIENT_COLOR,
diffuse_color: Union[float, Tuple[float]] = BASIC_DIFFUSE_COLOR,
specular_color: Union[float, Tuple[float]] = 0.0,
output_srgb: bool = True,
device: torch.device = torch.device("cuda"),
**kwargs,
):
super().__init__(**kwargs)
assert isinstance(volume, BoundingBoxVolume), "cannot sample points in unknown volume"
assert (nerstf is not None) ^ (sdf is not None and tf is not None)
self.sdf = sdf
self.tf = tf
self.nerstf = nerstf
self.void = void
self.volume = volume
self.grid_size = grid_size
self.n_coarse_samples = n_coarse_samples
self.n_fine_samples = n_fine_samples
self.importance_sampling_options = AttrDict(importance_sampling_options or {})
self.separate_shared_samples = separate_shared_samples
self.texture_channels = texture_channels
self.channel_scale = to_torch(channel_scale).to(device)
self.ambient_color = ambient_color
self.diffuse_color = diffuse_color
self.specular_color = specular_color
self.output_srgb = output_srgb
self.device = device
self.to(device)
def _query(
self,
query: Query,
params: AttrDict[str, torch.Tensor],
options: AttrDict[str, Any],
) -> AttrDict:
no_dir_query = query.copy()
no_dir_query.direction = None
if options.get("rendering_mode", "stf") == "stf":
assert query.direction is None
if self.nerstf is not None:
sdf = tf = self.nerstf(
query,
params=subdict(params, "nerstf"),
options=options,
)
else:
sdf = self.sdf(no_dir_query, params=subdict(params, "sdf"), options=options)
tf = self.tf(query, params=subdict(params, "tf"), options=options)
return AttrDict(
density=sdf.density,
signed_distance=sdf.signed_distance,
channels=tf.channels,
aux_losses=dict(),
)
def render_rays(
self,
batch: AttrDict,
params: Optional[Dict] = None,
options: Optional[AttrDict] = None,
) -> AttrDict:
"""
:param batch: has
- rays: [batch_size x ... x 2 x 3] specify the origin and direction of each ray.
:param options: Optional[Dict]
"""
params = self.update(params)
options = AttrDict() if options is None else AttrDict(options)
# Necessary to tell the TF to use specific NeRF channels.
options.rendering_mode = "nerf"
model = partial(self._query, params=params, options=options)
# First, render rays with coarse, stratified samples.
options.nerf_level = "coarse"
parts = [
RayVolumeIntegral(
model=model,
volume=self.volume,
sampler=StratifiedRaySampler(),
n_samples=self.n_coarse_samples,
),
]
coarse_results, samplers, coarse_raw_outputs = render_rays(
batch.rays,
parts,
self.void,
shared=not self.separate_shared_samples,
render_with_direction=options.render_with_direction,
importance_sampling_options=self.importance_sampling_options,
)
# Then, render with additional importance-weighted ray samples.
options.nerf_level = "fine"
parts = [
RayVolumeIntegral(
model=model,
volume=self.volume,
sampler=samplers[0],
n_samples=self.n_fine_samples,
),
]
fine_results, _, raw_outputs = render_rays(
batch.rays,
parts,
self.void,
shared=not self.separate_shared_samples,
prev_raw_outputs=coarse_raw_outputs,
render_with_direction=options.render_with_direction,
)
raw = raw_outputs[0]
aux_losses = fine_results.output.aux_losses.copy()
if self.separate_shared_samples:
for key, val in coarse_results.output.aux_losses.items():
aux_losses[key + "_coarse"] = val
channels = fine_results.output.channels
shape = [1] * (channels.ndim - 1) + [len(self.texture_channels)]
channels = channels * self.channel_scale.view(*shape)
res = AttrDict(
channels=channels,
transmittance=fine_results.transmittance,
raw_signed_distance=raw.signed_distance,
raw_density=raw.density,
distances=fine_results.output.distances,
t0=fine_results.volume_range.t0,
t1=fine_results.volume_range.t1,
intersected=fine_results.volume_range.intersected,
aux_losses=aux_losses,
)
if self.separate_shared_samples:
res.update(
dict(
channels_coarse=(
coarse_results.output.channels * self.channel_scale.view(*shape)
),
distances_coarse=coarse_results.output.distances,
transmittance_coarse=coarse_results.transmittance,
)
)
return res
def render_views(
self,
batch: AttrDict,
params: Optional[Dict] = None,
options: Optional[AttrDict] = None,
) -> AttrDict:
"""
Returns a backproppable rendering of a view
:param batch: contains either ["poses", "camera"], or ["cameras"]. Can
optionally contain any of ["height", "width", "query_batch_size"]
:param params: Meta parameters
contains rendering_mode in ["stf", "nerf"]
:param options: controls checkpointing, caching, and rendering.
Can provide a `rendering_mode` in ["stf", "nerf"]
"""
params = self.update(params)
options = AttrDict() if options is None else AttrDict(options)
if options.cache is None:
created_cache = True
options.cache = AttrDict()
else:
created_cache = False
rendering_mode = options.get("rendering_mode", "stf")
if rendering_mode == "nerf":
output = render_views_from_rays(
self.render_rays,
batch,
params=params,
options=options,
device=self.device,
)
elif rendering_mode == "stf":
sdf_fn = tf_fn = nerstf_fn = None
if self.nerstf is not None:
nerstf_fn = partial(
self.nerstf.forward_batched,
params=subdict(params, "nerstf"),
options=options,
)
else:
sdf_fn = partial(
self.sdf.forward_batched,
params=subdict(params, "sdf"),
options=options,
)
tf_fn = partial(
self.tf.forward_batched,
params=subdict(params, "tf"),
options=options,
)
output = render_views_from_stf(
batch,
options,
sdf_fn=sdf_fn,
tf_fn=tf_fn,
nerstf_fn=nerstf_fn,
volume=self.volume,
grid_size=self.grid_size,
channel_scale=self.channel_scale,
texture_channels=self.texture_channels,
ambient_color=self.ambient_color,
diffuse_color=self.diffuse_color,
specular_color=self.specular_color,
output_srgb=self.output_srgb,
device=self.device,
)
else:
raise NotImplementedError
if created_cache:
del options["cache"]
return output
def get_signed_distance(
self,
query: Query,
params: Dict[str, torch.Tensor],
options: AttrDict[str, Any],
) -> torch.Tensor:
if self.sdf is not None:
return self.sdf(query, params=subdict(params, "sdf"), options=options).signed_distance
assert self.nerstf is not None
return self.nerstf(query, params=subdict(params, "nerstf"), options=options).signed_distance
def get_texture(
self,
query: Query,
params: Dict[str, torch.Tensor],
options: AttrDict[str, Any],
) -> torch.Tensor:
if self.tf is not None:
return self.tf(query, params=subdict(params, "tf"), options=options).channels
assert self.nerstf is not None
return self.nerstf(query, params=subdict(params, "nerstf"), options=options).channels