tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

write eval combined images for diagnostics (#3070)

Browse Source
This commit is contained in:
J.Y 2024-04-11 14:37:02 -04:00 committed by GitHub
parent eba72db22a
commit 4714ae75f7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 9 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
nerfstudio/pipelines/base_pipeline.py
View File

@ -26,6 +26,7 @@ from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Uni
import torch
import torch.distributed as dist
import torchvision.utils as vutils
from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn
from torch import nn
from torch.cuda.amp.grad_scaler import GradScaler
@ -361,6 +362,8 @@ class VanillaPipeline(Pipeline):
metrics_dict_list = []
assert isinstance(self.datamanager, (VanillaDataManager, ParallelDataManager, FullImageDatamanager))
num_images = len(self.datamanager.fixed_indices_eval_dataloader)
if output_path is not None:
output_path.mkdir(exist_ok=True, parents=True)
with Progress(
TextColumn("[progress.description]{task.description}"),
BarColumn(),
@ -369,15 +372,18 @@ class VanillaPipeline(Pipeline):
transient=True,
) as progress:
task = progress.add_task("[green]Evaluating all eval images...", total=num_images)
idx = 0
for camera, batch in self.datamanager.fixed_indices_eval_dataloader:
# time this the following line
inner_start = time()
outputs = self.model.get_outputs_for_camera(camera=camera)
height, width = camera.height, camera.width
num_rays = height * width
metrics_dict, _ = self.model.get_image_metrics_and_images(outputs, batch)
metrics_dict, image_dict = self.model.get_image_metrics_and_images(outputs, batch)
if output_path is not None:
raise NotImplementedError("Saving images is not implemented yet")
for key in image_dict.keys():
image = image_dict[key] # [H, W, C] order
vutils.save_image(image.permute(2, 0, 1).cpu(), output_path / f"eval_{key}_{idx:04d}.png")
assert "num_rays_per_sec" not in metrics_dict
metrics_dict["num_rays_per_sec"] = (num_rays / (time() - inner_start)).item()
@ -386,6 +392,7 @@ class VanillaPipeline(Pipeline):
metrics_dict[fps_str] = (metrics_dict["num_rays_per_sec"] / (height * width)).item()
metrics_dict_list.append(metrics_dict)
progress.advance(task)
idx = idx + 1
# average the metrics list
metrics_dict = {}
for key in metrics_dict_list[0].keys():