mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
384 lines
11 KiB
Plaintext
384 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "3c00934d-0f14-4c1a-a563-94d627b377e5",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Notebook to visualize a dataset\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b15dfdc1-5054-40bc-9631-4c7426eecfd9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from notebook_imports import *\n",
|
|
"\n",
|
|
"import torch\n",
|
|
"from torch.utils.data import DataLoader\n",
|
|
"\n",
|
|
"from mattport.nerf.dataset.image_dataset import ImageDataset, collate_batch\n",
|
|
"from mattport.nerf.dataset.collate import CollateIterDataset, collate_batch_size_one\n",
|
|
"from mattport.nerf.dataset.utils import get_dataset_inputs\n",
|
|
"from mattport.nerf.field_modules.ray_generator import RayGenerator\n",
|
|
"from mattport.nerf.field_modules.scene_bounds_collider import SceneBoundsCollider\n",
|
|
"from mattport.structures.rays import RayBundle\n",
|
|
"from mattport.utils.io import get_absolute_path\n",
|
|
"from mattport.viewer.plotly import get_line_segments_from_lines\n",
|
|
"from mattport.structures.cameras import get_camera_model\n",
|
|
"\n",
|
|
"\n",
|
|
"from hydra import compose, initialize\n",
|
|
"from omegaconf import open_dict\n",
|
|
"import pprint\n",
|
|
"from tqdm import tqdm\n",
|
|
"import random"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "012bf657",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import goat\n",
|
|
"\n",
|
|
"vis = goat.vis_utils.get_vis()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "71f46c79",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"goat.vis_utils.show_box_test(vis)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "7fe7478d",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"with initialize(config_path=\"../configs\"):\n",
|
|
" cfg = compose(config_name=\"default.yaml\")\n",
|
|
" # cfg = compose(config_name=\"default.yaml\", overrides=[\"data/dataset=friends_TBBT-big_living_room\"])"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0b010534",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"dataset_inputs = get_dataset_inputs(**cfg.data.dataset, split=\"train\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "e95073a9",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# import goat.threedvis.meshcat.geometry as g\n",
|
|
"# import numpy as np\n",
|
|
"# point_cloud = dataset_inputs.point_cloud\n",
|
|
"# size = 0.1\n",
|
|
"# vis[\"point_cloud\"].set_object(\n",
|
|
"# g.PointCloud(position=np.transpose(np.array(point_cloud.xyz)),\n",
|
|
"# color=np.transpose(np.array(point_cloud.rgb)) / 255.0, size=size)\n",
|
|
"# )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "aac97d4a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# show dataset inputs with three.js\n",
|
|
"# goat.vis_utils.show_box_test(vis)\n",
|
|
"num_images_to_show = 200\n",
|
|
"for idx in tqdm(list(range(len(dataset_inputs.image_filenames)))[:num_images_to_show]):\n",
|
|
" image = imageio.imread(dataset_inputs.image_filenames[idx])[:, :, :3]\n",
|
|
"\n",
|
|
" num_intrinsics_params = len(dataset_inputs.intrinsics[idx])\n",
|
|
" camera_class = get_camera_model(num_intrinsics_params)\n",
|
|
" camera_to_world = dataset_inputs.camera_to_world[idx]\n",
|
|
" camera = camera_class(*dataset_inputs.intrinsics[idx].tolist(), camera_to_world=camera_to_world)\n",
|
|
" K = camera.get_intrinsics_matrix()\n",
|
|
"\n",
|
|
" name = \"train/{:06d}\".format(idx)\n",
|
|
" # height = 0.5\n",
|
|
" displayed_focal_length = 0.5\n",
|
|
" goat.vis_utils.draw_camera_frustum(\n",
|
|
" vis,\n",
|
|
" name=name,\n",
|
|
" image=image,\n",
|
|
" K=K.numpy(),\n",
|
|
" pose=camera_to_world.double().numpy(),\n",
|
|
" height=3.0,\n",
|
|
" realistic=False,\n",
|
|
" displayed_focal_length=0.0,\n",
|
|
" )"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "001bdf3b-60a0-4bb2-9f88-ae36f2036c80",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"image_dataset = ImageDataset(\n",
|
|
" image_filenames=dataset_inputs.image_filenames, downscale_factor=dataset_inputs.downscale_factor\n",
|
|
")\n",
|
|
"iter_dataset = CollateIterDataset(\n",
|
|
" image_dataset,\n",
|
|
" collate_fn=lambda batch_list: collate_batch(\n",
|
|
" batch_list, cfg.data.dataloader.num_rays_per_batch, keep_full_image=True\n",
|
|
" ),\n",
|
|
" num_samples_to_collate=cfg.data.dataloader.num_images_to_sample_from,\n",
|
|
" num_times_to_repeat=cfg.data.dataloader.num_times_to_repeat_images,\n",
|
|
")\n",
|
|
"dataloader = DataLoader(\n",
|
|
" iter_dataset,\n",
|
|
" batch_size=1,\n",
|
|
" num_workers=cfg.data.dataloader.num_workers,\n",
|
|
" collate_fn=collate_batch_size_one,\n",
|
|
" pin_memory=True,\n",
|
|
")\n",
|
|
"dataloader_iter = iter(dataloader)\n",
|
|
"\n",
|
|
"ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)\n",
|
|
"\n",
|
|
"num_batches = 10\n",
|
|
"for _ in tqdm(range(num_batches)):\n",
|
|
" batch = next(dataloader_iter)\n",
|
|
" ray_bundle = ray_generator.forward(batch[\"indices\"])\n",
|
|
" break"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "19b85061-a1b0-458b-9d00-96cde6f37dca",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"media.show_image(image_dataset.get_image(10))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "b90f4a78-4b8b-4615-a127-7d7579874765",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"def show_batch(batch):\n",
|
|
" # set the color of the sampled rays\n",
|
|
" c, y, x = [i.flatten() for i in torch.split(batch[\"local_indices\"], 1, dim=-1)]\n",
|
|
" batch[\"image\"][c, y, x] = 1.0\n",
|
|
"\n",
|
|
" # batch[\"image\"] is num_images, h, w, 3\n",
|
|
" images = torch.split(batch[\"image\"], 1, dim=0)\n",
|
|
" image_list = [image[0] for image in images]\n",
|
|
" image = torch.cat(image_list, dim=1) # cat along the width dimension\n",
|
|
"\n",
|
|
" # the white pixels are rays\n",
|
|
" media.show_image((image * 255).to(torch.uint8))\n",
|
|
"\n",
|
|
"\n",
|
|
"def sample_and_show_batch():\n",
|
|
" batch = next(iter(dataloader))\n",
|
|
" show_batch(batch)\n",
|
|
" return batch"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "09e625f7-fa1f-4b3f-acd4-bddfcdf7cd9f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"batch = sample_and_show_batch()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fd8cf48a-823e-4395-99ca-eda4c2026440",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"ray_indices = batch[\"indices\"]\n",
|
|
"ray_bundle = ray_generator(ray_indices)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "646107c0-f979-408a-941f-98b2ad3e07f7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"skip = 1\n",
|
|
"size = 8\n",
|
|
"data = []\n",
|
|
"data += [\n",
|
|
" go.Scatter3d(\n",
|
|
" x=ray_generator.camera_to_world[::skip, 0, 3],\n",
|
|
" y=ray_generator.camera_to_world[::skip, 1, 3],\n",
|
|
" z=ray_generator.camera_to_world[::skip, 2, 3],\n",
|
|
" mode=\"markers\",\n",
|
|
" name=\"origins\",\n",
|
|
" marker=dict(color=\"rgba(0, 0, 0, 1)\", size=size),\n",
|
|
" )\n",
|
|
"]"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "fc627f05",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"scene_bounds_collider = SceneBoundsCollider(dataset_inputs.scene_bounds)\n",
|
|
"intersected_ray_bundle = scene_bounds_collider(ray_bundle)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9ca1b3c3",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# intersected_ray_bundle.nears"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "c73ba9bd-2902-4c46-b592-bc87ac8c71f6",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# near_plane = cfg.graph.network.near_plane\n",
|
|
"# far_plane = cfg.graph.network.far_plane\n",
|
|
"lines = torch.stack(\n",
|
|
" [\n",
|
|
" intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.nears,\n",
|
|
" intersected_ray_bundle.origins + intersected_ray_bundle.directions * intersected_ray_bundle.fars,\n",
|
|
" ],\n",
|
|
" dim=1,\n",
|
|
").tolist() # (num_rays, 2, 3)\n",
|
|
"lines = torch.tensor(random.sample(lines, k=100))"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "991df6a2-efe3-4412-ba41-9a244e2e90cc",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"data += get_line_segments_from_lines(lines)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0b2404d8-ca9e-4630-9f67-7c68d19458cb",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"layout = go.Layout(\n",
|
|
" autosize=False,\n",
|
|
" width=1000,\n",
|
|
" height=1000,\n",
|
|
" margin=go.layout.Margin(l=50, r=50, b=100, t=100, pad=4),\n",
|
|
" scene=go.layout.Scene(\n",
|
|
" aspectmode=\"data\",\n",
|
|
" camera=dict(up=dict(x=0, y=0, z=1), center=dict(x=0, y=0, z=0), eye=dict(x=1.25, y=1.25, z=1.25)),\n",
|
|
" ),\n",
|
|
")\n",
|
|
"fig = go.Figure(data=data, layout=layout)\n",
|
|
"fig.show()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "7aa506fe-66a1-4284-ade4-8c90fde5476a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f259edb0-4126-4930-99f4-b5574dbb2eea",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"fig.write_image(\"temp.png\")\n",
|
|
"plotly_image = imageio.imread(\"temp.png\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "d1a11772-4692-4800-b446-cfdbe4ab92fa",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "722f5370-3b0f-4ad1-87a1-df479ae85b60",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": []
|
|
}
|
|
],
|
|
"metadata": {
|
|
"interpreter": {
|
|
"hash": "4a749acc7d255f078aee52e0584cc77b3cb5aaed1b3a7407ec4262c1bf6cb526"
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "mattport",
|
|
"language": "python",
|
|
"name": "mattport"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.8.13"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|