tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

viewer fix to be compatible with rgb vs rgb_fine/coarse (#439)

Browse Source 
small viewer fixes
This commit is contained in:
evonneng 2022-09-15 10:01:25 -07:00 committed by GitHub
parent 8973039e65
commit 1d6592b760
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 58 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

80
nerfactory/viewer/server/viewer_utils.py
View File

@ -51,6 +51,8 @@ class OutputTypes(str, enum.Enum):
INIT = "init"
RGB = "rgb"
RGB_FINE = "rgb_fine"
ACCUMULATION = "accumulation"
ACCUMULATION_FINE = "accumulation_fine"
class ColormapTypes(str, enum.Enum):
@ -238,7 +240,7 @@ class VisualizerState:
self.camera_moving = False
self.prev_camera_timestamp = 0
self.outputs_set = False
self.output_list = None
def init_scene(self, dataset: InputDataset, start_train=True) -> None:
"""Draw some images and the scene aabb in the viewer.
@ -323,7 +325,7 @@ class VisualizerState:
batches_per_sec = train_rays_per_sec / num_rays_per_batch
num_steps = int(1 / self.static_fps * batches_per_sec)
num_steps = max(int(1 / self.static_fps * batches_per_sec), 1)
else:
num_steps = 1
@ -376,40 +378,48 @@ class VisualizerState:
stuff_colors: is only set if colormap is for semantics. Defaults to None.
eps: epsilon to handle floating point comparisons
"""
if self.output_list:
reformatted_output = self._process_invalid_output(self.prev_output_type)
# default for rgb images
if self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].shape[-1] == 3:
return outputs[self.prev_output_type]
if self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].shape[-1] == 3:
return outputs[reformatted_output]
# rendering depth outputs
if self.prev_colormap_type == ColormapTypes.DEPTH or (
self.prev_colormap_type == ColormapTypes.DEFAULT
and outputs[self.prev_output_type].dtype == torch.float
and (torch.max(outputs[self.prev_output_type]) - 1.0) > eps # handle floating point arithmetic
and outputs[reformatted_output].dtype == torch.float
and (torch.max(outputs[reformatted_output]) - 1.0) > eps # handle floating point arithmetic
):
accumulation_str = (
OutputTypes.ACCUMULATION
if OutputTypes.ACCUMULATION in self.output_list
else OutputTypes.ACCUMULATION_FINE
)
return visualization.apply_depth_colormap(
outputs[self.prev_output_type], accumulation=outputs["accumulation"]
outputs[reformatted_output], accumulation=outputs[accumulation_str]
)
# rendering accumulation outputs
if self.prev_colormap_type == ColormapTypes.TURBO or (
self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].dtype == torch.float
self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.float
):
return visualization.apply_colormap(outputs[self.prev_output_type])
return visualization.apply_colormap(outputs[reformatted_output])
# rendering semantic outputs
if self.prev_colormap_type == ColormapTypes.SEMANTIC or (
self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].dtype == torch.int
self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.int
):
logits = outputs[self.prev_output_type]
logits = outputs[reformatted_output]
labels = torch.argmax(torch.nn.functional.softmax(logits, dim=-1), dim=-1) # type: ignore
assert stuff_colors is not None
return stuff_colors[labels]
# rendering boolean outputs
if self.prev_colormap_type == ColormapTypes.BOOLEAN or (
self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[self.prev_output_type].dtype == torch.bool
self.prev_colormap_type == ColormapTypes.DEFAULT and outputs[reformatted_output].dtype == torch.bool
):
return visualization.apply_boolean_colormap(outputs[self.prev_output_type])
return visualization.apply_boolean_colormap(outputs[reformatted_output])
raise NotImplementedError
@ -421,20 +431,24 @@ class VisualizerState:
stuff_colors: is only set if colormap is for semantics. Defaults to None.
eps: epsilon to handle floating point comparisons
"""
if not self.outputs_set:
self.vis["renderingState/output_options"].write(list(outputs.keys()))
self.outputs_set = True
# gross hack to get the image key, depending on which keys the graph uses
if self.prev_output_type == OutputTypes.INIT:
self.prev_output_type = OutputTypes.RGB if OutputTypes.RGB in outputs else OutputTypes.RGB_FINE
if self.output_list is None:
self.output_list = list(outputs.keys())
viewer_output_list = list(np.copy(self.output_list))
# remapping rgb_fine -> rgb for all cases just so that we dont have 2 of them in the options
if OutputTypes.RGB_FINE in self.output_list:
viewer_output_list.remove(OutputTypes.RGB_FINE)
viewer_output_list.insert(0, OutputTypes.RGB)
self.vis["renderingState/output_options"].write(viewer_output_list)
reformatted_output = self._process_invalid_output(self.prev_output_type)
# re-register colormaps and send to viewer
if self.output_type_changed or self.prev_colormap_type == ColormapTypes.INIT:
self.prev_colormap_type = ColormapTypes.DEFAULT
colormap_options = [ColormapTypes.DEFAULT]
if (
outputs[self.prev_output_type].shape[-1] != 3
and outputs[self.prev_output_type].dtype == torch.float
and (torch.max(outputs[self.prev_output_type]) - 1.0) <= eps # handle floating point arithmetic
outputs[reformatted_output].shape[-1] != 3
and outputs[reformatted_output].dtype == torch.float
and (torch.max(outputs[reformatted_output]) - 1.0) <= eps # handle floating point arithmetic
):
# accumulation can also include depth
colormap_options.extend(["depth"])
@ -518,6 +532,28 @@ class VisualizerState:
return image_height
def _process_invalid_output(self, output_type: str) -> str:
"""Check to see whether we are in the corner case of RGB; if still invalid, throw error
Returns correct string mapping given improperly formatted output_type.
Args:
output_type: reformatted output type
"""
if output_type == OutputTypes.INIT:
output_type = OutputTypes.RGB
# check if rgb or rgb_fine should be the case TODO: add other checks here
attempted_output_type = output_type
if output_type not in self.output_list and output_type == OutputTypes.RGB:
output_type = OutputTypes.RGB_FINE
# check if output_type is not in list
if output_type not in self.output_list:
assert (
NotImplementedError
), f"Output {attempted_output_type} not in list. Tried to reformat as {output_type} but still not found."
return output_type
@profiler.time_function
def _render_image_in_viewer(self, camera_object, graph: Model, is_training: bool) -> None:
"""