{ "cells": [ { "cell_type": "markdown", "id": "3c00934d-0f14-4c1a-a563-94d627b377e5", "metadata": {}, "source": [ "# Notebook to visualize a dataset\n", "\n", "Sanity check to ensure data is properly read\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", "\n", "from pyrad.data.dataloader import TrainDataloader\n", "from pyrad.data.utils import get_dataset_inputs_from_dataset_config\n", "from pyrad.graphs.modules.ray_generator import RayGenerator\n", "from pyrad.graphs.modules.scene_colliders import AABBBoxCollider\n", "from pyrad.utils.plotly import get_line_segments_from_lines\n", "from pyrad.utils.misc import instantiate_from_dict_config\n", "\n", "\n", "from hydra import compose, initialize\n", "from tqdm import tqdm\n", "import random" ] }, { "cell_type": "markdown", "id": "2cb2fd5f", "metadata": {}, "source": [ "## Load dataset" ] }, { "cell_type": "code", "execution_count": 25, "id": "7fe7478d", "metadata": {}, "outputs": [], "source": [ "CONFIG_NAME = \"graph_default.yaml\"\n", "DEVICE = \"cpu\"\n", "\n", "with initialize(version_base=None, config_path=\"../configs\"):\n", " config = compose(config_name=CONFIG_NAME)\n", "\n", "dataset_inputs = get_dataset_inputs_from_dataset_config(**config.data.dataset_inputs_train, split=\"train\")\n", "\n", "# ImageDataset\n", "image_dataset_train = instantiate_from_dict_config(config.data.image_dataset_train, **dataset_inputs.as_dict())\n", "# ImageSampler\n", "image_sampler_train = instantiate_from_dict_config(\n", " config.data.dataloader_train.image_sampler, dataset=image_dataset_train, device=DEVICE\n", ")\n", "# PixelSampler\n", "pixel_sampler_train = instantiate_from_dict_config(config.data.dataloader_train.pixel_sampler, keep_full_image=True)\n", "\n", "ray_generator = RayGenerator(dataset_inputs.intrinsics, dataset_inputs.camera_to_world)\n", "\n", "dataloader_train = TrainDataloader(image_sampler_train, pixel_sampler_train)\n", "iter_dataloader_train = iter(dataloader_train)\n", "\n", "\n", "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[\"full_image\"][c, y, x] = 0.0\n", " batch[\"full_image\"][c, y, x, 0] = 1.0\n", "\n", " # batch[\"image\"] is num_images, h, w, 3\n", " images = torch.split(batch[\"full_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 pink pixels are rays\n", " media.show_image((image * 255).to(torch.uint8))\n", "\n", "\n", "def sample_and_show_batch():\n", " ray_indices, batch = next(iter_dataloader_train)\n", " show_batch(batch)\n", " return batch" ] }, { "cell_type": "markdown", "id": "39349ea0", "metadata": {}, "source": [ "## Visualize Image\n" ] }, { "cell_type": "code", "execution_count": 28, "id": "19b85061-a1b0-458b-9d00-96cde6f37dca", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Example image:\n" ] }, { "data": { "text/html": [ "