data_gen2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Testing of ResNet+MLP where we see how well it performs when trained using a NeRF but tested on 'actual' camera data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some useful settings for interactive work\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functions.nerf import *\n",
    "from pathlib import Path\n",
    "import os\n",
    "import torchvision\n",
    "import numpy as np\n",
    "from scipy.spatial.transform import Rotation as R\n",
    "import json\n",
    "from torchvision.models import ResNet18_Weights\n",
    "import torchvision.transforms as T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[01:55:36] </span>Caching <span style=\"color: #800080; text-decoration-color: #800080\">/</span> undistorting train images                                            <a href=\"file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">full_images_datamanager.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py#182\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">182</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[01:55:36]\u001b[0m\u001b[2;36m \u001b[0mCaching \u001b[35m/\u001b[0m undistorting train images                                            \u001b]8;id=151915;file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py\u001b\\\u001b[2mfull_images_datamanager.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=689249;file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py#182\u001b\\\u001b[2m182\u001b[0m\u001b]8;;\u001b\\\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "368139a802a343d884478d4d20ab8e4f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #7fbfbf; text-decoration-color: #7fbfbf\">[01:55:41] </span>Caching <span style=\"color: #800080; text-decoration-color: #800080\">/</span> undistorting eval images                                             <a href=\"file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">full_images_datamanager.py</span></a><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">:</span><a href=\"file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py#196\" target=\"_blank\"><span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">196</span></a>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[2;36m[01:55:41]\u001b[0m\u001b[2;36m \u001b[0mCaching \u001b[35m/\u001b[0m undistorting eval images                                             \u001b]8;id=639590;file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py\u001b\\\u001b[2mfull_images_datamanager.py\u001b[0m\u001b]8;;\u001b\\\u001b[2m:\u001b[0m\u001b]8;id=402257;file:///home/lowjunen/anaconda3/envs/sfti2-env/lib/python3.10/site-packages/nerfstudio/data/datamanagers/full_images_datamanager.py#196\u001b\\\u001b[2m196\u001b[0m\u001b]8;;\u001b\\\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "34131dd92d964a69a7ce91a4e12ac54c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Loading latest checkpoint from load_dir\n",
       "</pre>\n"
      ],
      "text/plain": [
       "Loading latest checkpoint from load_dir\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">✅ Done loading checkpoint from \n",
       "outputs/colmap_in_mocap/splatfacto/<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">2024</span>-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">04</span>-15_192132/nerfstudio_models/step-<span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">000029999.</span>ckpt\n",
       "</pre>\n"
      ],
      "text/plain": [
       "✅ Done loading checkpoint from \n",
       "outputs/colmap_in_mocap/splatfacto/\u001b[1;36m2024\u001b[0m-\u001b[1;36m04\u001b[0m-15_192132/nerfstudio_models/step-\u001b[1;36m000029999.\u001b[0mckpt\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "os.chdir('/home/lowjunen/StanfordMSL/Image2Pose/nerfs')\n",
    "gs_pipeline = NeRF(Path('outputs/colmap_in_mocap/splatfacto/2024-04-15_192132/config.yml'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "preprocess = ResNet18_Weights.DEFAULT.transforms()\n",
    "transform = T.ToPILImage()\n",
    "\n",
    "def drone2nerf_transform(pose):\n",
    "    pd = pose[:3]\n",
    "    qd = pose[3:]\n",
    "\n",
    "    T_b2f = np.eye(4)\n",
    "    T_b2f[:3,:3],T_b2f[:3,3] = R.from_quat(qd).as_matrix(),pd\n",
    "\n",
    "    # Flightroom Frame to NeRF world frame\n",
    "    T_f2n = np.eye(4)\n",
    "    T_f2n[0:3,0:3] = np.array([\n",
    "        [ 1.0, 0.0, 0.0],\n",
    "        [ 0.0,-1.0, 0.0],\n",
    "        [ 0.0, 0.0,-1.0]\n",
    "    ])\n",
    "\n",
    "    # Camera Frame to NeRF world frame\n",
    "    T_c2b = np.eye(4)\n",
    "    T_c2b[0:3,0:3] = np.array([\n",
    "        [ 0.0, 0.0,-1.0],\n",
    "        [ 1.0, 0.0, 0.0],\n",
    "        [ 0.0,-1.0, 0.0]\n",
    "    ])\n",
    "\n",
    "    # Get image\n",
    "    T_c2n = T_f2n@T_b2f@T_c2b\n",
    "    \n",
    "    return T_c2n\n",
    "\n",
    "def waypoints_to_poses(tFOkf:np.ndarray,hz:int=10):\n",
    "    Nkf = tFOkf.shape[1]\n",
    "\n",
    "    FOs = []\n",
    "    for i in range(Nkf-1):\n",
    "        N = int((tFOkf[0,i+1]-tFOkf[0,i])*hz)\n",
    "\n",
    "        fo = np.linspace(tFOkf[1:5,i],tFOkf[1:5,i+1],N+1).T\n",
    "        if i < Nkf-2:\n",
    "            fo = fo[:,0:-1]\n",
    "\n",
    "\n",
    "        FOs.append(fo)\n",
    "\n",
    "    FOs = np.concatenate(FOs,axis=1)\n",
    "\n",
    "    poses = np.zeros((7,FOs.shape[1]))\n",
    "    poses[0:3,:] = FOs[0:3,:]\n",
    "    for i in range(FOs.shape[1]):\n",
    "        poses[3:,i] = R.from_euler('z',FOs[3,i]).as_quat()\n",
    "\n",
    "    return poses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "tFOkf = np.array([\n",
    "    [ 0.00, 1.00, 4.00, 5.00],\n",
    "    [-2.00, 0.00, 3.00, 5.00],\n",
    "    [-0.40,-0.40, 2.00, 2.00],\n",
    "    [-1.00,-1.30,-1.30,-1.30],\n",
    "    [ 0.00, 0.70, 0.30, 0.00]\n",
    "])\n",
    "\n",
    "poses = waypoints_to_poses(tFOkf,50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Desired Pose (in flightroom frame)\n",
    "frames = []\n",
    "idx = 0\n",
    "for j in range(15):\n",
    "    for i in range(poses.shape[1]):\n",
    "        # Generate Transform\n",
    "        T_c2n = drone2nerf_transform(poses[:,i])\n",
    "\n",
    "        # Generate noise\n",
    "        w_q = np.random.uniform(-0.1,0.1,4)\n",
    "        w_p = np.random.uniform(-0.3,0.3,3)\n",
    "\n",
    "        # Add noise\n",
    "        qi = R.from_matrix(T_c2n[0:3,0:3]).as_quat()+w_q\n",
    "        qi = qi/np.linalg.norm(qi)\n",
    "        ti = T_c2n[0:3,3]+w_p    \n",
    "        \n",
    "        T_c2n[0:3,0:3] = R.from_quat(qi).as_matrix()\n",
    "        T_c2n[0:3,3] = ti\n",
    "\n",
    "        # Render Image\n",
    "        pose = torch.tensor(T_c2n[0:3,:]).float()\n",
    "        outputs = gs_pipeline.render(pose)\n",
    "\n",
    "        img = torch.movedim(outputs['rgb'],2,0)\n",
    "        img_rnet = preprocess(transform(img))\n",
    "\n",
    "        file_name = 'data_'+str(idx).zfill(5)+'.png'\n",
    "        file_path = 'resn_data/'+file_name\n",
    "        torchvision.utils.save_image(img_rnet, '../data/'+file_path)\n",
    "        frames.append({\"file_path\": file_path,\"transform_matrix\":pose.tolist()})\n",
    "\n",
    "        # Update index\n",
    "        idx += 1\n",
    "\n",
    "with open('../data/resn_data/transforms.json', 'w') as json_file:\n",
    "    json.dump({\"frames\":frames}, json_file, indent=4)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "i2p-env",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}