mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
Improve depth vis
This commit is contained in:
Matt Tancik
parent
470812fab6
commit
b4e45d6fc0
@ -16,7 +16,7 @@ from mattport.nerf.field_modules.mlp import MLP
|
||||
from mattport.nerf.field_modules.ray_generator import RayGenerator
|
||||
from mattport.nerf.graph.base import Graph
|
||||
from mattport.nerf.loss import MSELoss
|
||||
from mattport.nerf.renderers import AccumulationRenderer, DisparityRenderer, RGBRenderer
|
||||
from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer
|
||||
from mattport.nerf.sampler import PDFSampler, UniformSampler
|
||||
from mattport.structures import colors
|
||||
from mattport.structures.rays import RaySamples
|
||||
@ -114,7 +114,7 @@ class NeRFGraph(Graph):
|
||||
# renderers
|
||||
self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
|
||||
self.renderer_accumulation = AccumulationRenderer()
|
||||
self.renderer_disparity = DisparityRenderer()
|
||||
self.renderer_depth = DepthRenderer()
|
||||
|
||||
# losses
|
||||
self.rgb_loss = MSELoss()
|
||||
@ -148,7 +148,7 @@ class NeRFGraph(Graph):
|
||||
weights=coarse_weights,
|
||||
) # RendererOutputs
|
||||
coarse_renderer_accumulation = self.renderer_accumulation(coarse_weights) # RendererOutputs
|
||||
coarse_renderer_disparity = self.renderer_disparity(coarse_weights, uniform_ray_samples.ts)
|
||||
coarse_renderer_depth = self.renderer_depth(coarse_weights, uniform_ray_samples.ts)
|
||||
|
||||
# fine network:
|
||||
pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights) # RaySamples
|
||||
@ -161,9 +161,21 @@ class NeRFGraph(Graph):
|
||||
weights=fine_weights,
|
||||
) # RendererOutputs
|
||||
fine_renderer_accumulation = self.renderer_accumulation(fine_weights) # RendererOutputs
|
||||
fine_renderer_disparity = self.renderer_disparity(fine_weights, pdf_ray_samples.ts)
|
||||
fine_renderer_depth = self.renderer_depth(fine_weights, pdf_ray_samples.ts)
|
||||
|
||||
# TODO refactor this into "vis" section. Doesn't need to be run during training.
|
||||
coarse_renderer_depth = visualization.apply_depth_colormap(
|
||||
coarse_renderer_depth.depth,
|
||||
accumulation=coarse_renderer_accumulation.accumulation,
|
||||
near_plane=self.near_plane,
|
||||
far_plane=self.far_plane,
|
||||
)
|
||||
fine_renderer_depth = visualization.apply_depth_colormap(
|
||||
fine_renderer_depth.depth,
|
||||
accumulation=fine_renderer_accumulation.accumulation,
|
||||
near_plane=self.near_plane,
|
||||
far_plane=self.far_plane,
|
||||
)
|
||||
coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation)
|
||||
fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation)
|
||||
|
||||
@ -173,8 +185,8 @@ class NeRFGraph(Graph):
|
||||
"rgb_fine": fine_renderer_outputs.rgb,
|
||||
"accumulation_coarse": coarse_renderer_accumulation,
|
||||
"accumulation_fine": fine_renderer_accumulation,
|
||||
"disparity_coarse": coarse_renderer_disparity.disparity,
|
||||
"disparity_fine": fine_renderer_disparity.disparity,
|
||||
"depth_coarse": coarse_renderer_depth,
|
||||
"depth_fine": fine_renderer_depth,
|
||||
}
|
||||
return outputs
|
||||
|
||||
|
||||
@ -16,7 +16,7 @@ class RendererOutputs:
|
||||
rgb: TensorType["num_rays", 3] = None
|
||||
density: TensorType["num_rays", 1] = None
|
||||
accumulation: TensorType["num_rays", 1] = None
|
||||
disparity: TensorType["num_rays", 1] = None
|
||||
depth: TensorType["num_rays", 1] = None
|
||||
|
||||
|
||||
class RGBRenderer(nn.Module):
|
||||
@ -82,7 +82,7 @@ class AccumulationRenderer(nn.Module):
|
||||
return renderer_outputs
|
||||
|
||||
|
||||
class DisparityRenderer(nn.Module):
|
||||
class DepthRenderer(nn.Module):
|
||||
"""Calcualte depth along ray."""
|
||||
|
||||
def __init__(self, method: str = "expected") -> None:
|
||||
@ -103,14 +103,15 @@ class DisparityRenderer(nn.Module):
|
||||
ts (TensorType[..., "num_samples"]): Sample locations along rays
|
||||
|
||||
Returns:
|
||||
RendererOutputs: Outputs with disparity values.
|
||||
RendererOutputs: Outputs with depth values.
|
||||
"""
|
||||
|
||||
if self.method == "expected":
|
||||
depth = torch.sum(weights * ts, dim=-1)
|
||||
eps = 1e-10
|
||||
disparity = 1.0 / torch.max(eps * torch.ones_like(depth), depth / torch.sum(weights, -1))
|
||||
depth = torch.sum(weights * ts, dim=-1) / (torch.sum(weights, -1) + eps)
|
||||
|
||||
return RendererOutputs(disparity=disparity)
|
||||
depth = torch.clip(depth, ts[..., 0], ts[..., -1])
|
||||
|
||||
return RendererOutputs(depth=depth[..., None])
|
||||
|
||||
raise NotImplementedError(f"Method {self.method} not implemented")
|
||||
|
||||
@ -259,8 +259,8 @@ class Trainer:
|
||||
rgb_fine = []
|
||||
accumulation_coarse = []
|
||||
accumulation_fine = []
|
||||
disparity_coarse = []
|
||||
disparity_fine = []
|
||||
depth_coarse = []
|
||||
depth_fine = []
|
||||
for i in range(0, num_rays, chunk_size):
|
||||
ray_indices = all_ray_indices[i : i + chunk_size].to(self.device)
|
||||
graph_outputs = self.graph(ray_indices)
|
||||
@ -268,14 +268,14 @@ class Trainer:
|
||||
rgb_fine.append(graph_outputs["rgb_fine"])
|
||||
accumulation_coarse.append(graph_outputs["accumulation_coarse"])
|
||||
accumulation_fine.append(graph_outputs["accumulation_fine"])
|
||||
disparity_coarse.append(graph_outputs["disparity_coarse"])
|
||||
disparity_fine.append(graph_outputs["disparity_fine"])
|
||||
depth_coarse.append(graph_outputs["depth_coarse"])
|
||||
depth_fine.append(graph_outputs["depth_fine"])
|
||||
rgb_coarse = torch.cat(rgb_coarse).view(image_height, image_width, 3).detach().cpu()
|
||||
rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu()
|
||||
accumulation_coarse = torch.cat(accumulation_coarse).view(image_height, image_width, 3).detach().cpu()
|
||||
accumulation_fine = torch.cat(accumulation_fine).view(image_height, image_width, 3).detach().cpu()
|
||||
disparity_coarse = torch.cat(disparity_coarse).view(image_height, image_width, 1).detach().cpu()
|
||||
disparity_fine = torch.cat(disparity_fine).view(image_height, image_width, 1).detach().cpu()
|
||||
depth_coarse = torch.cat(depth_coarse).view(image_height, image_width, 3).detach().cpu()
|
||||
depth_fine = torch.cat(depth_fine).view(image_height, image_width, 3).detach().cpu()
|
||||
|
||||
combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1)
|
||||
writer.write_event(
|
||||
@ -287,10 +287,8 @@ class Trainer:
|
||||
{"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_accumulation"}
|
||||
)
|
||||
|
||||
combined_image = torch.cat([disparity_coarse, disparity_fine], dim=1)
|
||||
writer.write_event(
|
||||
{"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_disparity"}
|
||||
)
|
||||
combined_image = torch.cat([depth_coarse, depth_fine], dim=1)
|
||||
writer.write_event({"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_depth"})
|
||||
|
||||
coarse_psnr = get_psnr(image, rgb_coarse)
|
||||
writer.write_event(
|
||||
|
||||
@ -27,7 +27,6 @@ def write_event(args: Dict[str, Any]):
|
||||
assert comms.is_main_process(), "Writing out with process other than main"
|
||||
for writer in EVENT_WRITERS:
|
||||
if "x" in args:
|
||||
print(args["x"].shape)
|
||||
writer.write_image(**args)
|
||||
elif "scalar" in args:
|
||||
writer.write_scalar(**args)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user