dense_geometry.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.renderer.panopli_tensoRF_renderer import TensoRFRenderer\n",
    "from model.radiance_field.tensoRF import TensorVMSplit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Meta Platforms, Inc. All Rights Reserved\n",
    "import sys\n",
    "import random\n",
    "import omegaconf\n",
    "import torch\n",
    "from pathlib import Path\n",
    "from PIL import Image\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.utils import make_grid\n",
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "from dataset import PanopLiDataset, create_segmentation_data_panopli\n",
    "from model.radiance_field.tensoRF import TensorVMSplit\n",
    "from model.renderer.panopli_tensoRF_renderer import TensoRFRenderer\n",
    "from trainer import visualize_panoptic_outputs\n",
    "from util.camera import distance_to_depth\n",
    "from util.misc import get_parameters_from_state_dict\n",
    "\n",
    "from util.distinct_colors import DistinctColors\n",
    "from dataset.preprocessing.preprocess_scannet import get_thing_semantics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "cfg = omegaconf.OmegaConf.load(Path('/home/rozenberszki/project/panoptic-lifting/pretrained-examples/scannet_scene042302/checkpoints/epoch=9-step=678528.ckpt').parents[1] / \"config.yaml\")\n",
    "cfg.resume ='/home/rozenberszki/project/panoptic-lifting/pretrained-examples/scannet_scene042302/checkpoints/epoch=9-step=678528.ckpt'\n",
    "test_mode = False if len(sys.argv) == 2 else sys.argv[2] == \"True\"\n",
    "cfg.image_dim = [256, 384]\n",
    "if isinstance(cfg.image_dim, int):\n",
    "    cfg.image_dim = [cfg.image_dim, cfg.image_dim]\n",
    "config = cfg\n",
    "trajectory_name = \"trajectory_blender\"\n",
    "test_only = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "runs/scene0423_02_test_01170151_PanopLi_scannet042302_electrical-forest\n",
      "\n",
      "[00] aabb tensor([-1., -1., -1.,  1.,  1.,  1.])\n",
      "[00] grid size tensor([128, 128, 128])\n",
      "[00] units:  None\n",
      "[00] sampling step size:  tensor(0.0079)\n",
      "[00] sampling number:  440\n",
      "\n",
      "[00] aabb tensor([-1.0000, -0.9528, -0.0866,  1.0000,  1.0000,  0.3071])\n",
      "[00] grid size tensor([332, 324,  65])\n",
      "[00] units:  tensor(0.0079)\n",
      "[00] sampling step size:  tensor(0.0030)\n",
      "[00] sampling number:  929\n",
      "\n"
     ]
    }
   ],
   "source": [
    "output_dir = (\n",
    "    Path(\"runs\")\n",
    "    / f\"{Path(config.dataset_root).stem}_{trajectory_name if not test_only else 'test'}_{Path(config.experiment)}\"\n",
    ")\n",
    "print(output_dir)\n",
    "output_dir.mkdir(parents=True, exist_ok=True)\n",
    "device = torch.device(\"cuda:0\")\n",
    "test_set = PanopLiDataset(\n",
    "    Path(config.dataset_root),\n",
    "    \"test\",\n",
    "    (config.image_dim[0], config.image_dim[1]),\n",
    "    config.max_depth,\n",
    "    overfit=config.overfit,\n",
    "    semantics_dir=\"m2f_semantics\",\n",
    "    instance_dir=\"m2f_instance\",\n",
    "    instance_to_semantic_key=\"m2f_instance_to_semantic\",\n",
    "    create_seg_data_func=create_segmentation_data_panopli,\n",
    "    subsample_frames=config.subsample_frames,\n",
    ")\n",
    "H, W, alpha = config.image_dim[0], config.image_dim[1], 0.65\n",
    "# whether to render the test set or a predefined trajectory through the scene\n",
    "if test_only:\n",
    "    trajectory_set = test_set\n",
    "else:\n",
    "    trajectory_set = test_set.get_trajectory_set(trajectory_name, True)\n",
    "trajectory_loader = DataLoader(\n",
    "    trajectory_set, shuffle=False, num_workers=0, batch_size=1\n",
    ")\n",
    "checkpoint = torch.load(config.resume, map_location=\"cpu\")\n",
    "state_dict = checkpoint[\"state_dict\"]\n",
    "total_classes = len(test_set.segmentation_data.bg_classes) + len(\n",
    "    test_set.segmentation_data.fg_classes\n",
    ")\n",
    "output_mlp_semantics = (\n",
    "    torch.nn.Identity()\n",
    "    if config.semantic_weight_mode != \"softmax\"\n",
    "    else torch.nn.Softmax(dim=-1)\n",
    ")\n",
    "model = TensorVMSplit(\n",
    "    [config.min_grid_dim, config.min_grid_dim, config.min_grid_dim],\n",
    "    num_semantics_comps=(32, 32, 32),\n",
    "    num_semantic_classes=total_classes,\n",
    "    dim_feature_instance=config.max_instances,\n",
    "    output_mlp_semantics=output_mlp_semantics,\n",
    "    use_semantic_mlp=config.use_mlp_for_semantics,\n",
    ")\n",
    "renderer = TensoRFRenderer(\n",
    "    test_set.scene_bounds,\n",
    "    [config.min_grid_dim, config.min_grid_dim, config.min_grid_dim],\n",
    "    semantic_weight_mode=config.semantic_weight_mode,\n",
    ")\n",
    "renderer.load_state_dict(get_parameters_from_state_dict(state_dict, \"renderer\"))\n",
    "for epoch in config.grid_upscale_epochs[::-1]:\n",
    "    if checkpoint[\"epoch\"] >= epoch:\n",
    "        model.upsample_volume_grid(renderer.grid_dim)\n",
    "        renderer.update_step_size(renderer.grid_dim)\n",
    "        break\n",
    "model.load_state_dict(get_parameters_from_state_dict(state_dict, \"model\"))\n",
    "model = model.to(device)\n",
    "renderer = renderer.to(device)\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "grid_conv: 100%|██████████| 107/107 [00:01<00:00, 63.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "        # create a new model for background\n",
    "        chunk_size = int(2 ** 16)\n",
    "        alpha, dense_xyz = renderer.get_dense_alpha(model)\n",
    "        xyz_sampled = renderer.normalize_coordinates(dense_xyz)\n",
    "        outputs_semantic = []\n",
    "        outputs_instance = []\n",
    "        output_colors = []\n",
    "        for chunk in tqdm(torch.split(xyz_sampled.view(-1, 3), chunk_size), desc='grid_conv'):\n",
    "            outputs_semantic.append(model.render_semantic_mlp(None, model.compute_semantic_feature(chunk)).cpu())\n",
    "            outputs_instance.append(model.compute_instance_feature(chunk).cpu())\n",
    "            dir_0 = torch.zeros_like(chunk)\n",
    "            dir_a, dir_b, dir_c = dir_0.clone(), dir_0.clone(), dir_0.clone()\n",
    "            dir_a[:, 0] = 1\n",
    "            dir_b[:, 1] = 1\n",
    "            dir_c[:, 2] = 1\n",
    "            total = model.render_appearance_mlp(dir_a, model.compute_appearance_feature(chunk)).cpu() + model.render_appearance_mlp(dir_b, model.compute_appearance_feature(chunk)).cpu() + model.render_appearance_mlp(dir_c, model.compute_appearance_feature(chunk)).cpu()\n",
    "            total = total / 3\n",
    "            output_colors.append(total)\n",
    "        sem_labels = torch.cat(outputs_semantic, dim=0).reshape([xyz_sampled.shape[0], xyz_sampled.shape[1], xyz_sampled.shape[2], -1]).argmax(dim=-1).int().transpose(0, 2).contiguous()\n",
    "        colors = torch.cat(output_colors, dim=0).reshape([xyz_sampled.shape[0], xyz_sampled.shape[1], xyz_sampled.shape[2], -1]).float().transpose(0, 2).contiguous()\n",
    "        dense_xyz = dense_xyz.transpose(0, 2).contiguous()\n",
    "        alpha = alpha.clamp(0, 1).transpose(0, 2).contiguous()\n",
    "        alpha_thres = 0.01 ## no idea set by me, Julius\n",
    "        alpha[alpha >= alpha_thres] = 1\n",
    "        alpha[alpha < alpha_thres] = 0\n",
    "        mask = alpha > 0.5\n",
    "        valid_xyz = dense_xyz[mask]\n",
    "        distinct_colors = DistinctColors()\n",
    "        semantic_bg = torch.from_numpy(np.isin(sem_labels.cpu().numpy(), [i for i, x in enumerate(get_thing_semantics(\"extended\")) if not x])).to(sem_labels.device)\n",
    "        thing_semantics = sem_labels.clone()\n",
    "        thing_semantics[semantic_bg] = 0\n",
    "        colored_semantics = distinct_colors.get_color_fast_numpy(sem_labels.cpu().numpy().reshape(-1)).reshape(list(sem_labels.shape) + [3])\n",
    "        print('')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1421244, 3])"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "valid_xyz.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([65, 324, 332]), (65, 324, 332, 3))"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "alpha.shape, colored_semantics.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'normalized_mc' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[51], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mnormalized_mc\u001b[49m(alpha\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m2\u001b[39m), torch\u001b[38;5;241m.\u001b[39mfrom_numpy(colored_semantics)\u001b[38;5;241m.\u001b[39mtranspose(\u001b[38;5;241m0\u001b[39m, \u001b[38;5;241m2\u001b[39m), renderer\u001b[38;5;241m.\u001b[39mbbox_aabb)\u001b[38;5;241m.\u001b[39mexport(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124msemantics.obj\u001b[39m\u001b[38;5;124m'\u001b[39m)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'normalized_mc' is not defined"
     ]
    }
   ],
   "source": [
    "normalized_mc(alpha.transpose(0, 2), torch.from_numpy(colored_semantics).transpose(0, 2), renderer.bbox_aabb).export('semantics.obj')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([598242, 3]), (65, 324, 332, 3))"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "valid_xyz.shape, colored_semantics.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "import trimesh\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "# Assuming your tensor is named 'points'\n",
    "points = valid_xyz.detach().cpu() # Replace with your actual tensor\n",
    "# Convert the tensor to a numpy array\n",
    "points_np = points.numpy()\n",
    "\n",
    "# Create a trimesh object from the numpy array\n",
    "mesh = trimesh.Trimesh(vertices=points_np)\n",
    "\n",
    "# Save the mesh as a .ply file\n",
    "mesh.export('output_scannet.ply', file_type='ply')\n",
    "print('')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "gt = trimesh.load(\"/home/rozenberszki/scene0423_02/scene0423_02_vh_clean_2.ply\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(60071, 3)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gt.vertices.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "panoptic",
   "language": "python",
   "name": "python3"
  },
  "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}