tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 18:44:21 +08:00
Code

Merge branch 'master' of github.com:ethanweber/radiance

Browse Source
This commit is contained in:
Matt Tancik 2022-05-31 14:08:47 -07:00
commit ee1dcdd197
7 changed files with 40 additions and 82 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
radiance/nerf/graph/instant_ngp.py
View File

@ -130,12 +130,12 @@ class NGPGraph(Graph):
combined_acc = torch.cat([acc], dim=1)
combined_depth = torch.cat([depth], dim=1)
writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth")
writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
depth = visualization.apply_depth_colormap(outputs["depth_occupancy_grid"])
writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth_occupancy_grid")
writer.put_image(name=f"depth_occupancy_grid/image_idx_{image_idx}", image=combined_depth, step=step)
# Switch images from [H, W, C] to [1, C, H, W] for metrics computations
image = torch.moveaxis(image, -1, 0)[None, ...]
@ -145,8 +145,8 @@ class NGPGraph(Graph):
ssim = self.ssim(image, rgb)
lpips = self.lpips(image, rgb)
writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(psnr), step=step, group="psnr")
writer.put_scalar(name=f"val_{image_idx}", scalar=float(ssim), step=step, group="ssim")
writer.put_scalar(name=f"val_{image_idx}", scalar=float(lpips), step=step, group="lpips")
writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(psnr), step=step)
writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(ssim), step=step)
writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(lpips), step=step)
return psnr.item()

2
radiance/nerf/graph/nerfw.py
View File

@ -210,4 +210,4 @@ class NerfWGraph(NeRFGraph):
row2 = torch.cat([depth_fine, depth_fine_static, depth_coarse], dim=-2)
combined_image = torch.cat([row0, row1, row2], dim=-3)
writer.put_image(name="image_idx_{image_idx}-nerfw", image=combined_image, step=step, group="img")
writer.put_image(name="img/image_idx_{image_idx}-nerfw", image=combined_image, step=step)

2
radiance/nerf/graph/semantic_nerf.py
View File

@ -155,4 +155,4 @@ class SemanticNerfGraph(NeRFGraph):
self.stuff_classes
)
semantic_labels_image = visualization.apply_colormap(semantic_labels[..., None])
writer.put_image(name="image_idx_{image_idx}", image=semantic_labels_image, step=step, group="semantics")
writer.put_image(name="semantics/image_idx_{image_idx}", image=semantic_labels_image, step=step)

16
radiance/nerf/graph/vanilla_nerf.py
View File

@ -152,10 +152,10 @@ class NeRFGraph(Graph):
combined_depth = torch.cat([depth_coarse, depth_fine], dim=1)
mask = visualization.apply_depth_colormap(mask[..., None])
writer.put_image(name=f"image_idx_{image_idx}", image=combined_rgb, step=step, group="img")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_acc, step=step, group="accumulation")
writer.put_image(name=f"image_idx_{image_idx}", image=combined_depth, step=step, group="depth")
writer.put_image(name=f"image_idx_{image_idx}", image=mask, step=step, group="mask")
writer.put_image(name=f"img/image_idx_{image_idx}", image=combined_rgb, step=step)
writer.put_image(name=f"accumulation/image_idx_{image_idx}", image=combined_acc, step=step)
writer.put_image(name=f"depth/image_idx_{image_idx}", image=combined_depth, step=step)
writer.put_image(name=f"mask/image_idx_{image_idx}", image=mask, step=step)
# Switch images from [H, W, C] to [1, C, H, W] for metrics computations
image = torch.moveaxis(image, -1, 0)[None, ...]
@ -167,10 +167,10 @@ class NeRFGraph(Graph):
fine_ssim = self.ssim(image, rgb_fine)
fine_lpips = self.lpips(image, rgb_fine)
writer.put_scalar(name=f"val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step, group="psnr")
writer.put_scalar(name=f"val_{image_idx}-fine", scalar=float(fine_psnr), step=step, group="psnr")
writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_ssim), step=step, group="ssim")
writer.put_scalar(name=f"val_{image_idx}", scalar=float(fine_lpips), step=step, group="lpips")
writer.put_scalar(name=f"psnr/val_{image_idx}-coarse", scalar=float(coarse_psnr), step=step)
writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step)
writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step)
writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step)
writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step)

2
radiance/nerf/optimizers.py
View File

@ -93,4 +93,4 @@ class Optimizers:
scheduler.step()
# TODO(ethan): clean this up. why is there indexing into a list?
lr = scheduler.get_last_lr()[0]
writer.put_scalar(name=f"{param_group_name}", scalar=lr, step=step, group="learning_rate")
writer.put_scalar(name=f"learning_rate/{param_group_name}", scalar=lr, step=step)

2
radiance/nerf/trainer.py
View File

@ -230,7 +230,7 @@ class Trainer:
)
if step != 0 and step % self.config.logging.steps_per_log == 0:
writer.put_dict(name="loss_dict", scalar_dict=loss_dict, step=step, group="Loss", prefix="train-")
writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step)
# TODO: add the learning rates to tensorboard/logging
if step != 0 and self.config.graph.steps_per_save and step % self.config.graph.steps_per_save == 0:
self.save_checkpoint(self.config.graph.model_dir, step)

84
radiance/utils/writer.py
View File

@ -45,40 +45,27 @@ class EventType(enum.Enum):
@check_main_thread
def put_image(name, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None):
def put_image(name, image: TensorType["H", "W", "C"], step: int):
"""Setter function to place images into the queue to be written out"""
if isinstance(name, EventName):
name = name.value
EVENT_STORAGE.append(
{"name": name, "write_type": EventType.IMAGE, "event": image, "step": step, "group": group, "prefix": prefix}
)
EVENT_STORAGE.append({"name": name, "write_type": EventType.IMAGE, "event": image, "step": step})
@check_main_thread
def put_scalar(name: str, scalar: float, step: int, group: str = None, prefix: str = None):
def put_scalar(name: str, scalar: float, step: int):
"""Setter function to place scalars into the queue to be written out"""
if isinstance(name, EventName):
name = name.value
EVENT_STORAGE.append(
{"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step, "group": group, "prefix": prefix}
)
EVENT_STORAGE.append({"name": name, "write_type": EventType.SCALAR, "event": scalar, "step": step})
@check_main_thread
def put_dict(name: str, scalar_dict: float, step: int, group: str = None, prefix: str = None):
def put_dict(name: str, scalar_dict: float, step: int):
"""Setter function to place a dictionary of scalars into the queue to be written out"""
EVENT_STORAGE.append(
{
"name": name,
"write_type": EventType.DICT,
"event": scalar_dict,
"step": step,
"group": group,
"prefix": prefix,
}
)
EVENT_STORAGE.append({"name": name, "write_type": EventType.DICT, "event": scalar_dict, "step": step})
@check_main_thread
@ -87,8 +74,6 @@ def put_time(
start_time: float,
end_time: float,
step: int,
group: str = None,
prefix: str = None,
avg_over_iters: bool = False,
avg_over_batch: int = None,
update_eta: bool = False,
@ -116,16 +101,16 @@ def put_time(
curr_buffer.pop(0)
curr_buffer.append(val)
curr_avg = sum(curr_buffer) / len(curr_buffer)
put_scalar(name, curr_avg, step, group, prefix)
put_scalar(name, curr_avg, step)
GLOBAL_BUFFER["events"][name] = {"buffer": curr_buffer, "avg": curr_avg}
else:
put_scalar(name, val, step, group, prefix)
put_scalar(name, val, step)
if update_eta:
## NOTE: eta should be called with avg train iteration time
remain_iter = GLOBAL_BUFFER["max_iter"] - step
remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"]
put_scalar(EventName.ETA, remain_time, step, group, prefix)
put_scalar(EventName.ETA, remain_time, step)
@check_main_thread
@ -135,11 +120,11 @@ def write_out_storage():
for event in EVENT_STORAGE:
write_func = getattr(writer, event["write_type"].value)
if event["write_type"] == EventType.DICT:
write_func(event["event"], event["step"], event["group"], event["prefix"])
write_func(event["event"], event["step"])
if isinstance(writer, LocalWriter):
continue
else:
write_func(event["name"], event["event"], event["step"], event["group"], event["prefix"])
write_func(event["name"], event["event"], event["step"])
if isinstance(writer, LocalWriter):
break
EVENT_STORAGE.clear()
@ -165,16 +150,6 @@ def setup_event_writers(config: DictConfig) -> None:
GLOBAL_BUFFER["events"] = {}
def get_tensorboard_name(name: str, group: str = None, prefix: str = None):
"""Returns a string for tensorboard with an optional group and prefix.
Where tensorboard_name has the form `group/prefix-name`.
"""
group_string = f"{group}/" if group else ""
prefix_string = f"{prefix}" if prefix else ""
tensorboard_name = f"{group_string}{prefix_string}{name}"
return tensorboard_name
class Writer:
"""Writer class"""
@ -182,46 +157,36 @@ class Writer:
self.save_dir = save_dir
@abstractmethod
def write_image(
self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None
) -> None:
def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None:
"""_summary_
Args:
name (str): data identifier
image (TensorType["H", "W", "C"]): rendered image to write
step (int): the time step to log
group (str): the group e.g., "Loss", "Accuracy", "Time"
prefix (str): the prefix e.g., "train-", "test-"
"""
raise NotImplementedError
@abstractmethod
def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
def write_scalar(self, name: str, scalar: float, step: int) -> None:
"""Required method to write a single scalar value to the logger
Args:
name (str): data identifier
step (int): the time step to log
group (str): the group e.g., "Loss", "Accuracy", "Time"
prefix (str): the prefix e.g., "train-", "test-"
"""
raise NotImplementedError
@check_main_thread
def write_scalar_dict(
self, scalar_dict: Dict[str, float], step: int, group: str = None, prefix: str = None
) -> None:
def write_scalar_dict(self, scalar_dict: Dict[str, float], step: int) -> None:
"""Function that writes out all scalars from a given dictionary to the logger
Args:
scalar_dict (dict): dictionary containing all scalar values with key names and quantities
step (int): the time step to log
group (str): the group e.g., "Loss", "Accuracy", "Time"
prefix (str): the prefix e.g., "train-", "test-"
"""
for name, scalar in scalar_dict.items():
self.write_scalar(name, scalar, step, group=group, prefix=prefix)
self.write_scalar(name, scalar, step)
@decorate_all([check_main_thread])
@ -232,9 +197,7 @@ class TensorboardWriter(Writer):
super().__init__(save_dir)
self.tb_writer = SummaryWriter(log_dir=self.save_dir)
def write_image(
self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None
) -> None:
def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None:
"""_summary_
Args:
@ -242,20 +205,17 @@ class TensorboardWriter(Writer):
image (TensorType["H", "W", "C"]): rendered image to write
"""
image = to8b(image)
tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix)
self.tb_writer.add_image(tensorboard_name, image, step, dataformats="HWC")
self.tb_writer.add_image(name, image, step, dataformats="HWC")
def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
def write_scalar(self, name: str, scalar: float, step: int) -> None:
"""Tensorboard method to write a single scalar value to the logger
Args:
name (str): data identifier
x (float): x value to write
y (float): y value to write
group (str)): a prefix to group tensorboard scalars
"""
tensorboard_name = get_tensorboard_name(name, group=group, prefix=prefix)
self.tb_writer.add_scalar(tensorboard_name, scalar, step)
self.tb_writer.add_scalar(name, scalar, step)
def _cursorup(x: int):
@ -303,15 +263,13 @@ class LocalWriter(Writer):
self.keys = set()
self.past_mssgs = ["", ""]
def write_image(
self, name: str, image: TensorType["H", "W", "C"], step: int, group: str = None, prefix: str = None
) -> None:
def write_image(self, name: str, image: TensorType["H", "W", "C"], step: int) -> None:
if name in self.stats_to_track:
image = to8b(image)
image_path = os.path.join(self.save_dir, f"{name}.jpg")
imageio.imwrite(image_path, np.uint8(image.cpu().numpy() * 255.0))
def write_scalar(self, name: str, scalar: float, step: int, group: str = None, prefix: str = None) -> None:
def write_scalar(self, name: str, scalar: float, step: int) -> None:
if step > 0:
latest_map, new_key = self._consolidate_events()
self._update_header(latest_map, new_key)