nerfstudio/notebooks/visualize_viewer.ipynb

Notebook to visualize the viewer functionality

In [ ]:

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "4"

In [ ]:

from notebook_imports import *

import time
import torch
from torch.utils.data import DataLoader

from pyrad.data.dataloader import TrainDataloader
from pyrad.data.image_dataset import ImageDataset, PanopticImageDataset
from pyrad.data.image_sampler import CacheImageSampler
from pyrad.data.pixel_sampler import PixelSampler
from pyrad.data.utils import DatasetInputs, get_dataset_inputs_from_dataset_config
from pyrad.graphs.modules.ray_generator import RayGenerator
from pyrad.graphs.modules.scene_colliders import SceneBoundsCollider, AABBBoxCollider
from pyrad.cameras.cameras import get_camera
from pyrad.cameras.camera_paths import InterpolatedCameraPath
from pyrad.cameras.rays import RayBundle
from pyrad.utils.io import get_absolute_path
from pyrad.utils.plotly import get_line_segments_from_lines
from pyrad.cameras.cameras import get_camera_model
from pyrad.utils.misc import get_dict_to_torch, instantiate_from_dict_config
from pyrad.engine.trainer import Trainer
from pyrad.utils.config import Config, setup_config

# the new import
from pyrad.viewer.backend import vis_utils


from hydra import compose, initialize
from omegaconf import open_dict
import pprint
from tqdm import tqdm
import random

import umsgpack
from pyrad.viewer.backend.utils import get_chunks, get_intrinsics_matrix_and_camera_to_world_h

In [ ]:

# establish the connection to the tcp backend, to talk to the visualizer
# zmq_url should match the output of `python run_zmq_server.py`
vis = vis_utils.get_vis(zmq_url="tcp://127.0.0.1:6000")

In [ ]:

vis.delete()

# draws a red box in the scene
vis_utils.show_box_test(vis)

In [ ]:

with initialize(version_base=None, config_path="../configs"):
    config = compose(
        config_name="graph_default.yaml",
        overrides=[
            "trainer.resume_train.load_dir=/shared/ethanweber/pyrad/outputs/blender_lego/vanilla_nerf/2022-07-04_221907/checkpoints",
            "trainer.resume_train.load_step=4000",
        ],
    )
    config = setup_config(config)

In [ ]:

trainer = Trainer(config, local_rank=0, world_size=1)
trainer.setup()
trainer.draw_scene_in_viewer()

In [ ]:

aabb = trainer.dataset_inputs_train.scene_bounds.aabb
vis_utils.draw_aabb(trainer.vis, aabb, name="dataset_inputs_train/scene_bounds/aabb")

In [ ]:

In [ ]:

outputs = trainer.render_image_in_viewer()

In [ ]:

trainer.vis.delete()

In [ ]:

while True:
    outputs = trainer.render_image_in_viewer()
    image = outputs["rgb_fine"].cpu().numpy().astype("float64") * 255
    trainer.vis["/Cameras/Main Camera"].set_image(image)

In [ ]:

image.shape

In [ ]:

dataset_inputs = get_dataset_inputs_from_dataset_config(**config.data.dataset_inputs_train, split="train")

In [ ]:

# ImageDataset
image_dataset_train = instantiate_from_dict_config(config.data.image_dataset_train, **dataset_inputs.as_dict())

In [ ]:

# move a camera around
num_cameras = len(dataset_inputs.intrinsics)
intrinsics = dataset_inputs.intrinsics
camera_to_world = dataset_inputs.camera_to_world
idx0, idx1 = random.sample(range(num_cameras), k=2)
camera_a = get_camera(intrinsics[idx0], camera_to_world[idx0], None)
camera_b = get_camera(intrinsics[idx1], camera_to_world[idx1], None)

num_steps = 100
fps = 30
estimated_time = num_steps / fps
print("estimated_time:", estimated_time)

camera_path = InterpolatedCameraPath(camera_a, camera_b)
cameras = camera_path.get_path(steps=num_steps)

start = time.time()
for camera in cameras:
    vis_utils.set_camera(vis, camera)
    time.sleep(1 / fps)
time_elapsed = time.time() - start
print("time_elapsed:", time_elapsed)

In [ ]: