tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 10:34:29 +08:00
Code

when camera_type is equirectangular, add spherical metadata to rendered video (#1151)

Browse Source 
* when camera_type is equirectangular, add spherical metadata to rendered video

* remove unnecessary CONSOLE.print

* refactor to make pylint happier

* get os.stat size properly; document camera_type argument

* write video on the fly to avoid holding all images in memory

Note: video render width and height were unspecified for spiral trajectory; this uses values for Nerfacto

* fix images output

* doc fix
This commit is contained in:
machenmusik 2022-12-21 00:32:42 -05:00 committed by GitHub
parent 4a513825c4
commit 0e1ea37308
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 137 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

166
scripts/render.py
View File

@ -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(
"""<rdf:SphericalVideo
xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'
xmlns:GSpherical='http://ns.google.com/videos/1.0/spherical/'>
<GSpherical:ProjectionType>equirectangular</GSpherical:ProjectionType>
<GSpherical:Spherical>True</GSpherical:Spherical>
<GSpherical:Stitched>True</GSpherical:Stitched>
<GSpherical:StitchingSoftware>nerfstudio</GSpherical:StitchingSoftware>
</rdf:SphericalVideo>""",
"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,
)