From 2c4fd33cc7cc48d78d3d5c365e56059c23278437 Mon Sep 17 00:00:00 2001 From: "Frederik Warburg frwa@dtu.dk" Date: Wed, 26 Jul 2023 12:43:18 +0000 Subject: [PATCH] allow to save per frame results with gt --- nerfstudio/scripts/render.py | 86 +++++++++++++++++++++--------------- 1 file changed, 51 insertions(+), 35 deletions(-) diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 60378616d..c7ec88793 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -20,8 +20,8 @@ from __future__ import annotations import json import os -import struct import shutil +import struct import sys from contextlib import ExitStack from dataclasses import dataclass, field @@ -35,23 +35,12 @@ import tyro from jaxtyping import Float from rich import box, style from rich.panel import Panel -from rich.progress import ( - BarColumn, - Progress, - TaskProgressColumn, - TextColumn, - TimeElapsedColumn, - TimeRemainingColumn, -) +from rich.progress import BarColumn, Progress, TaskProgressColumn, TextColumn, TimeElapsedColumn, TimeRemainingColumn from rich.table import Table from torch import Tensor from typing_extensions import Annotated -from nerfstudio.cameras.camera_paths import ( - get_interpolated_camera_path, - get_path_from_json, - get_spiral_path, -) +from nerfstudio.cameras.camera_paths import get_interpolated_camera_path, get_path_from_json, get_spiral_path from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager from nerfstudio.data.scene_box import SceneBox @@ -74,6 +63,7 @@ def _render_trajectory_video( output_format: Literal["images", "video"] = "video", image_format: Literal["jpeg", "png"] = "jpeg", jpeg_quality: int = 100, + concatenate_outputs: bool = False, colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions(), ) -> None: """Helper function to create a video of the spiral trajectory. @@ -87,6 +77,7 @@ def _render_trajectory_video( rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. seconds: Length of output video. output_format: How to save output data. + concatenate_outputs: Whether to concatenate the outputs images. colormap_options: Options for colormap. """ CONSOLE.print("[bold green]Creating trajectory " + output_format) @@ -139,7 +130,12 @@ def _render_trajectory_video( with torch.no_grad(): outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - render_image = [] + # also save ground truth if specified + if "rgb_gt" in rendered_output_names: + assert pipeline.datamanager.eval_dataset is not None + outputs["rgb_gt"] = pipeline.datamanager.eval_dataset.get_image(camera_idx).to(pipeline.device) + + render_images = [] for rendered_output_name in rendered_output_names: if rendered_output_name not in outputs: CONSOLE.rule("Error", style="red") @@ -157,27 +153,41 @@ def _render_trajectory_video( .cpu() .numpy() ) - render_image.append(output_image) - render_image = np.concatenate(render_image, axis=1) - if output_format == "images": - if image_format == "png": - media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image, fmt="png") - if image_format == "jpeg": - media.write_image( - output_image_dir / f"{camera_idx:05d}.jpg", render_image, fmt="jpeg", quality=jpeg_quality - ) - if output_format == "video": - if writer is None: - render_width = int(render_image.shape[1]) - render_height = int(render_image.shape[0]) - writer = stack.enter_context( - media.VideoWriter( - path=output_filename, - shape=(render_height, render_width), - fps=fps, + render_images.append(output_image) + + if concatenate_outputs: + render_images = [np.concatenate(render_images, axis=1)] + render_names = ["composite"] + else: + render_names = rendered_output_names + + for render_name, render_image in zip(render_names, render_images): + if output_format == "images": + os.makedirs(output_image_dir / render_name, exist_ok=True) + if image_format == "png": + media.write_image( + output_image_dir / render_name / f"{camera_idx:05d}.png", render_image, fmt="png" ) - ) - writer.add_image(render_image) + if image_format == "jpeg": + media.write_image( + output_image_dir / render_name / f"{camera_idx:05d}.jpg", + render_image, + fmt="jpeg", + quality=jpeg_quality, + ) + if output_format == "video": + if writer is None: + render_width = int(render_image.shape[1]) + render_height = int(render_image.shape[0]) + + writer = stack.enter_context( + media.VideoWriter( + path=output_image_dir / f"{render_name}.mp4", + shape=(render_height, render_width), + fps=fps, + ) + ) + writer.add_image(render_image) table = Table( title=None, @@ -315,6 +325,8 @@ class BaseRender: """Specifies number of rays per chunk during eval. If None, use the value in the config file.""" colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions() """Colormap options.""" + concatenate_outputs: bool = True + """Whether to concatenate the outputs images.""" @dataclass @@ -372,6 +384,7 @@ class RenderCameraPath(BaseRender): image_format=self.image_format, jpeg_quality=self.jpeg_quality, colormap_options=self.colormap_options, + concatenate_outputs=self.concatenate_outputs, ) if camera_path.camera_type[0] == CameraType.OMNIDIRECTIONALSTEREO_L.value: @@ -396,6 +409,7 @@ class RenderCameraPath(BaseRender): image_format=self.image_format, jpeg_quality=self.jpeg_quality, colormap_options=self.colormap_options, + concatenate_outputs=self.concatenate_outputs, ) # stack the left and right eye renders for final output @@ -471,6 +485,7 @@ class RenderInterpolated(BaseRender): output_format=self.output_format, image_format=self.image_format, colormap_options=self.colormap_options, + concatenate_outputs=self.concatenate_outputs, ) @@ -514,6 +529,7 @@ class SpiralRender(BaseRender): output_format=self.output_format, image_format=self.image_format, colormap_options=self.colormap_options, + concatenate_outputs=self.concatenate_outputs, )