mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-10 18:17:31 +08:00
cfff2bb65d
* key fixes * Update README.md * Update run_eval.py * Better error reporting * lint * Removed rgb fine change * Readme updates * Justify center
246 lines
9.2 KiB
Python
Executable File
246 lines
9.2 KiB
Python
Executable File
#!/usr/bin/env python
|
|
"""
|
|
eval.py
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import dataclasses
|
|
import json
|
|
import logging
|
|
import os
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Literal, Tuple, Union
|
|
|
|
import dcargs
|
|
import mediapy as media
|
|
import torch
|
|
import yaml
|
|
from rich.console import Console
|
|
from rich.progress import (
|
|
BarColumn,
|
|
Progress,
|
|
TaskProgressColumn,
|
|
TextColumn,
|
|
TimeRemainingColumn,
|
|
)
|
|
from typing_extensions import assert_never
|
|
|
|
from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path
|
|
from nerfstudio.cameras.cameras import Cameras
|
|
from nerfstudio.configs import base_config as cfg
|
|
from nerfstudio.pipelines.base_pipeline import Pipeline
|
|
from nerfstudio.utils.rich_utils import ItersPerSecColumn
|
|
|
|
console = Console(width=120)
|
|
|
|
logging.basicConfig(format="[%(filename)s:%(lineno)d] %(message)s", level=logging.INFO)
|
|
|
|
|
|
def _load_checkpoint(config: cfg.TrainerConfig, pipeline: Pipeline) -> Path:
|
|
"""Helper function to load checkpointed pipeline
|
|
|
|
Args:
|
|
config (DictConfig): Configuration of pipeline to load
|
|
pipeline (Pipeline): Pipeline instance of which to load weights
|
|
"""
|
|
assert config.load_dir is not None
|
|
if config.load_step is None:
|
|
console.print("Loading latest checkpoint from load_dir")
|
|
# NOTE: this is specific to the checkpoint name format
|
|
if not os.path.exists(config.load_dir):
|
|
console.rule("Error", style="red")
|
|
console.print(f"No checkpoint directory found at {config.load_dir}, ", justify="center")
|
|
console.print(
|
|
"Please make sure the checkpoint exists, they should be generated periodically during training",
|
|
justify="center",
|
|
)
|
|
sys.exit(1)
|
|
load_step = sorted(int(x[x.find("-") + 1 : x.find(".")]) for x in os.listdir(config.load_dir))[-1]
|
|
else:
|
|
load_step = config.load_step
|
|
load_path = config.load_dir / f"step-{load_step:09d}.ckpt"
|
|
assert load_path.exists(), f"Checkpoint {load_path} does not exist"
|
|
loaded_state = torch.load(load_path, map_location="cpu")
|
|
pipeline.load_pipeline(loaded_state["pipeline"])
|
|
console.print(f":white_check_mark: Done loading checkpoint from {load_path}")
|
|
return load_path
|
|
|
|
|
|
def _render_trajectory_video(
|
|
pipeline: Pipeline,
|
|
cameras: Cameras,
|
|
output_filename: Path,
|
|
rendered_output_name: str,
|
|
rendered_resolution_scaling_factor: float = 1.0,
|
|
seconds: float = 5.0,
|
|
) -> None:
|
|
"""Helper function to create a video of the spiral trajectory.
|
|
|
|
Args:
|
|
pipeline: Pipeline to evaluate with.
|
|
cameras: Cameras to render.
|
|
output_filename: Name of the output file.
|
|
rendered_output_name: Name of the renderer output to use.
|
|
rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution.
|
|
seconds: Number for the output video.
|
|
"""
|
|
console.print("[bold green]Creating trajectory video")
|
|
images = []
|
|
cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
|
|
|
|
progress = Progress(
|
|
TextColumn(":movie_camera: Rendering :movie_camera:"),
|
|
BarColumn(),
|
|
TaskProgressColumn(show_speed=True),
|
|
ItersPerSecColumn(suffix="fps"),
|
|
TimeRemainingColumn(elapsed_when_finished=True, compact=True),
|
|
)
|
|
with progress:
|
|
for camera_idx in progress.track(range(cameras.size), description=""):
|
|
camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx).to(pipeline.device)
|
|
with torch.no_grad():
|
|
outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
|
|
if rendered_output_name not in outputs:
|
|
console.rule("Error", style="red")
|
|
console.print(f"Could not find {rendered_output_name} in the model outputs", justify="center")
|
|
console.print(f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center")
|
|
sys.exit(1)
|
|
image = outputs[rendered_output_name].cpu().numpy()
|
|
images.append(image)
|
|
|
|
fps = len(images) / seconds
|
|
with console.status("[yellow]Saving video", spinner="bouncingBall"):
|
|
media.write_video(output_filename, images, fps=fps)
|
|
console.rule("[green] :tada: :tada: :tada: Success :tada: :tada: :tada:")
|
|
console.print(f"[green]Saved video to {output_filename}", justify="center")
|
|
|
|
|
|
def _eval_setup(config_path: Path) -> Tuple[cfg.Config, Pipeline, Path]:
|
|
"""Shared setup for loading a saved pipeline for evaluation.
|
|
|
|
Args:
|
|
config_path: Path to config YAML file.
|
|
|
|
Returns:
|
|
Loaded config, pipeline module, and corresponding checkpoint.
|
|
"""
|
|
# load save config
|
|
config = yaml.load(config_path.read_text(), Loader=yaml.Loader)
|
|
assert isinstance(config, cfg.Config)
|
|
|
|
# load checkpoints from wherever they were saved
|
|
# TODO: expose the ability to choose an arbitrary checkpoint
|
|
config.trainer.load_dir = config.get_checkpoint_dir()
|
|
config.pipeline.datamanager.eval_image_indices = None
|
|
|
|
# setup pipeline (which includes the DataManager)
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
pipeline = config.pipeline.setup(device=device, test_mode=True)
|
|
assert isinstance(pipeline, Pipeline)
|
|
pipeline.eval()
|
|
|
|
# load checkpointed information
|
|
checkpoint_path = _load_checkpoint(config.trainer, pipeline)
|
|
|
|
return config, pipeline, checkpoint_path
|
|
|
|
|
|
@dataclass
|
|
class ComputePSNR:
|
|
"""Load a checkpoint, compute some PSNR metrics, and save to a JSON."""
|
|
|
|
# Path to config YAML file.
|
|
load_config: Path
|
|
# Name of the output file.
|
|
output_path: Path = Path("output.json")
|
|
|
|
def main(self) -> None:
|
|
"""Main function."""
|
|
config, pipeline, checkpoint_path = _eval_setup(self.load_config)
|
|
|
|
assert self.output_path.suffix == ".json"
|
|
metrics_dict = pipeline.get_average_eval_image_metrics()
|
|
# save output to some file
|
|
self.output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
benchmark_info = {
|
|
"experiment_name": config.experiment_name,
|
|
"method_name": config.method_name,
|
|
"checkpoint": str(checkpoint_path),
|
|
"results": metrics_dict,
|
|
}
|
|
self.output_path.write_text(json.dumps(benchmark_info, indent=2), "utf8")
|
|
console.print(f"Saved results to: {self.output_path}")
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class RenderTrajectory:
|
|
"""Load a checkpoint, render a trajectory, and save to a video file."""
|
|
|
|
# Path to config YAML file.
|
|
load_config: Path
|
|
# Name of the renderer output to use. rgb, depth, etc.
|
|
rendered_output_name: str = "rgb"
|
|
# Trajectory to render.
|
|
traj: Literal["spiral", "interp", "filename"] = "spiral"
|
|
# Scaling factor to apply to the camera image resolution.
|
|
downscale_factor: int = 1
|
|
# Filename of the camera path to render.
|
|
camera_path_filename: Path = Path("camera_path.json")
|
|
# Name of the output file.
|
|
output_path: Path = Path("output.mp4")
|
|
# How long the video should be.
|
|
seconds: float = 5.0
|
|
|
|
def main(self) -> None:
|
|
"""Main function."""
|
|
_, pipeline, _ = _eval_setup(self.load_config)
|
|
|
|
seconds = self.seconds
|
|
|
|
# TODO(ethan): use camera information from parsing args
|
|
if self.traj == "spiral":
|
|
camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0)
|
|
# TODO(ethan): pass in the up direction of the camera
|
|
camera_path = get_spiral_path(camera_start, steps=30, radius=0.1)
|
|
elif self.traj == "interp":
|
|
# cameras_a = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0)
|
|
# cameras_b = pipeline.datamanager.eval_dataloader.get_camera(image_idx=10)
|
|
# camera_path = get_interpolated_camera_path(cameras, steps=30)
|
|
raise NotImplementedError("Interpolated camera path not implemented.")
|
|
elif self.traj == "filename":
|
|
with open(self.camera_path_filename, "r", encoding="utf-8") as f:
|
|
camera_path = json.load(f)
|
|
seconds = camera_path["seconds"]
|
|
camera_path = get_path_from_json(camera_path)
|
|
else:
|
|
assert_never(self.traj)
|
|
|
|
_render_trajectory_video(
|
|
pipeline,
|
|
camera_path,
|
|
output_filename=self.output_path,
|
|
rendered_output_name=self.rendered_output_name,
|
|
rendered_resolution_scaling_factor=1.0 / self.downscale_factor,
|
|
seconds=seconds,
|
|
)
|
|
|
|
|
|
def entrypoint():
|
|
"""Entrypoint for use with pyproject scripts."""
|
|
# A Union over dataclass types will create a subcommand for each type.
|
|
#
|
|
# TODO: it would make sense to split this script up into separate scripts.
|
|
# - To reduce duplicate code, some "upstream" refactor could also simplify the
|
|
# shared `_load_checkpoint()` and `_eval_setup()` helpers. The high-level
|
|
# operations implemented by each seem fairly universal; ideally the checkpoint
|
|
# loading logic, for example, would be the same as what's used for loading a
|
|
# checkpoint when resuming a training run.
|
|
dcargs.extras.set_accent_color("bright_yellow")
|
|
dcargs.cli(Union[ComputePSNR, RenderTrajectory]).main()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
entrypoint()
|