{ "cells": [ { "cell_type": "markdown", "id": "3c00934d-0f14-4c1a-a563-94d627b377e5", "metadata": {}, "source": [ "## Notebook to visualize a dataset\n" ] }, { "cell_type": "code", "execution_count": 1, "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_dict\n", "from mattport.nerf.field_modules.ray_generator import RayGenerator\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": 2, "id": "012bf657", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "The autoreload extension is already loaded. To reload it, use:\n", " %reload_ext autoreload\n" ] } ], "source": [ "import goat\n", "vis = goat.vis_utils.get_vis()" ] }, { "cell_type": "code", "execution_count": 3, "id": "71f46c79", "metadata": {}, "outputs": [], "source": [ "goat.vis_utils.show_box_test(vis)" ] }, { "cell_type": "code", "execution_count": 4, "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": 5, "id": "ad63e53a", "metadata": {}, "outputs": [], "source": [ "import os\n", "from mattport.nerf.dataset.colmap_utils import read_cameras_binary, read_images_binary, read_pointsTD_binary" ] }, { "cell_type": "code", "execution_count": 13, "id": "01c42863", "metadata": {}, "outputs": [], "source": [ "from mattport.nerf.dataset.friends import load_friends_data" ] }, { "cell_type": "code", "execution_count": null, "id": "b097e5eb", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 40, "id": "c5516443", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[2930.31056439 960. 544. ]\n", "[4466.94830951 1024. 574. ]\n" ] } ], "source": [ "dataset_inputs = load_friends_data(\"/shared/ethanweber/cartoons/data/colmap/smith_house_pair_manual_after\", include_point_cloud=True)" ] }, { "cell_type": "code", "execution_count": null, "id": "eba07414", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 24, "id": "8ff323ed", "metadata": {}, "outputs": [], "source": [ "imdata = read_images_binary(os.path.join(\"/shared/ethanweber/mattport/data/friends/TBBT-big_living_room\", \"colmap\", \"images.bin\"))" ] }, { "cell_type": "code", "execution_count": 41, "id": "10539dc6", "metadata": {}, "outputs": [], "source": [ "# cfg.data.dataset.data_directory = get_absolute_path(cfg.data.dataset.data_directory)\n", "# dataset_inputs = get_dataset_inputs_dict(**cfg.data.dataset, splits=[\"train\"])[\"train\"]" ] }, { "cell_type": "code", "execution_count": null, "id": "0b010534", "metadata": {}, "outputs": [], "source": [] }, { "cell_type": "code", "execution_count": 46, "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": 47, "id": "aac97d4a", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "100%|██████████| 2/2 [00:00<00:00, 4.15it/s]\n" ] } ], "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(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", " )\n", "\n" ] }, { "cell_type": "code", "execution_count": 8, "id": "001bdf3b-60a0-4bb2-9f88-ae36f2036c80", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ " 0%| | 0/10 [00:03, ?it/s]\n" ] } ], "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": 9, "id": "19b85061-a1b0-458b-9d00-96cde6f37dca", "metadata": {}, "outputs": [ { "data": { "text/html": [ "