diff --git a/scripts/render.py b/scripts/render.py index c8bb4fdae..5fa400303 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -5,7 +5,10 @@ render.py from __future__ import annotations import json +import os +import struct import sys +from contextlib import ExitStack from dataclasses import dataclass, field from pathlib import Path from typing import List, Optional @@ -25,7 +28,7 @@ from rich.progress import ( from typing_extensions import Literal, assert_never from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path -from nerfstudio.cameras.cameras import Cameras +from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.pipelines.base_pipeline import Pipeline from nerfstudio.utils import install_checks from nerfstudio.utils.eval_utils import eval_setup @@ -39,9 +42,12 @@ def _render_trajectory_video( cameras: Cameras, output_filename: Path, rendered_output_names: List[str], + render_width: int, + render_height: int, rendered_resolution_scaling_factor: float = 1.0, seconds: float = 5.0, output_format: Literal["images", "video"] = "video", + camera_type: CameraType = CameraType.PERSPECTIVE, ) -> None: """Helper function to create a video of the spiral trajectory. @@ -50,12 +56,14 @@ def _render_trajectory_video( cameras: Cameras to render. output_filename: Name of the output file. rendered_output_names: List of outputs to visualise. + render_width: Video width to render. + render_height: Video height to render. 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. + camera_type: Camera projection format type. """ - CONSOLE.print("[bold green]Creating trajectory video") - images = [] + CONSOLE.print("[bold green]Creating trajectory " + output_format) cameras.rescale_output_resolution(rendered_resolution_scaling_factor) cameras = cameras.to(pipeline.device) @@ -66,37 +74,121 @@ def _render_trajectory_video( ItersPerSecColumn(suffix="fps"), TimeRemainingColumn(elapsed_when_finished=True, compact=True), ) - output_image_dir = output_filename.parent / output_filename.stem if output_format == "images": + output_image_dir = output_filename.parent / output_filename.stem output_image_dir.mkdir(parents=True, exist_ok=True) - with progress: - for camera_idx in progress.track(range(cameras.size), description=""): - camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx) - with torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) - render_image = [] - for rendered_output_name in rendered_output_names: - 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) - output_image = outputs[rendered_output_name].cpu().numpy() - render_image.append(output_image) - render_image = np.concatenate(render_image, axis=1) - if output_format == "images": - media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image) - else: - images.append(render_image) - if output_format == "video": - fps = len(images) / seconds + fps = len(cameras) / seconds # make the folder if it doesn't exist output_filename.parent.mkdir(parents=True, exist_ok=True) - 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") + # NOTE: + # we could use ffmpeg_args "-movflags faststart" for progressive download, + # which would force moov atom into known position before mdat, + # but then we would have to move all of mdat to insert metadata atom + # (unless we reserve enough space to overwrite with our uuid tag, + # but we don't know how big the video file will be, so it's not certain!) + + with ExitStack() as stack: + writer = ( + stack.enter_context(media.VideoWriter(path=output_filename, shape=(render_height, render_width), fps=fps)) + if output_format == "video" + else None + ) + with progress: + for camera_idx in progress.track(range(cameras.size), description=""): + camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx) + with torch.no_grad(): + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + render_image = [] + for rendered_output_name in rendered_output_names: + 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) + output_image = outputs[rendered_output_name].cpu().numpy() + render_image.append(output_image) + render_image = np.concatenate(render_image, axis=1) + if output_format == "images": + media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image) + if output_format == "video": + writer.add_image(render_image) + + if output_format == "video": + if camera_type == CameraType.EQUIRECTANGULAR: + insert_spherical_metadata_into_file(output_filename) + + +def insert_spherical_metadata_into_file( + output_filename: Path, +) -> None: + """Inserts spherical metadata into MP4 video file in-place. + Args: + output_filename: Name of the (input and) output file. + """ + # NOTE: + # because we didn't use faststart, the moov atom will be at the end; + # to insert our metadata, we need to find (skip atoms until we get to) moov. + # we should have 0x00000020 ftyp, then 0x00000008 free, then variable mdat. + spherical_uuid = b"\xff\xcc\x82\x63\xf8\x55\x4a\x93\x88\x14\x58\x7a\x02\x52\x1f\xdd" + spherical_metadata = bytes( + """ +equirectangular +True +True +nerfstudio +""", + "utf-8", + ) + insert_size = len(spherical_metadata) + 8 + 16 + with open(output_filename, mode="r+b") as mp4file: + try: + # get file size + mp4file_size = os.stat(output_filename).st_size + + # find moov container (probably after ftyp, free, mdat) + while True: + pos = mp4file.tell() + size, tag = struct.unpack(">I4s", mp4file.read(8)) + if tag == b"moov": + break + mp4file.seek(pos + size) + # if moov isn't at end, bail + if pos + size != mp4file_size: + # TODO: to support faststart, rewrite all stco offsets + raise Exception("moov container not at end of file") + # go back and write inserted size + mp4file.seek(pos) + mp4file.write(struct.pack(">I", size + insert_size)) + # go inside moov + mp4file.seek(pos + 8) + # find trak container (probably after mvhd) + while True: + pos = mp4file.tell() + size, tag = struct.unpack(">I4s", mp4file.read(8)) + if tag == b"trak": + break + mp4file.seek(pos + size) + # go back and write inserted size + mp4file.seek(pos) + mp4file.write(struct.pack(">I", size + insert_size)) + # we need to read everything from end of trak to end of file in order to insert + # TODO: to support faststart, make more efficient (may load nearly all data) + mp4file.seek(pos + size) + rest_of_file = mp4file.read(mp4file_size - pos - size) + # go to end of trak (again) + mp4file.seek(pos + size) + # insert our uuid atom with spherical metadata + mp4file.write(struct.pack(">I4s16s", insert_size, b"uuid", spherical_uuid)) + mp4file.write(spherical_metadata) + # write rest of file + mp4file.write(rest_of_file) + finally: + mp4file.close() @dataclass @@ -138,11 +230,24 @@ class RenderTrajectory: if self.traj == "spiral": camera_start = pipeline.datamanager.eval_dataloader.get_camera(image_idx=0).flatten() # TODO(ethan): pass in the up direction of the camera + camera_type = CameraType.PERSPECTIVE + render_width = 952 + render_height = 736 camera_path = get_spiral_path(camera_start, steps=30, radius=0.1) 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"] + if "camera_type" not in camera_path: + camera_type = CameraType.PERSPECTIVE + elif camera_path["camera_type"] == "fisheye": + camera_type = CameraType.FISHEYE + elif camera_path["camera_type"] == "equirectangular": + camera_type = CameraType.EQUIRECTANGULAR + else: + camera_type = CameraType.PERSPECTIVE + render_width = camera_path["render_width"] + render_height = camera_path["render_height"] camera_path = get_path_from_json(camera_path) else: assert_never(self.traj) @@ -155,6 +260,9 @@ class RenderTrajectory: rendered_resolution_scaling_factor=1.0 / self.downscale_factor, seconds=seconds, output_format=self.output_format, + camera_type=camera_type, + render_width=render_width, + render_height=render_height, )