scene_recon/notebooks/Open3D_TSDF_integration_tensor.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b0ae020e",
   "metadata": {
    "code_folding": [
     13,
     25
    ]
   },
   "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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO - 2022-05-30 12:56:12,111 - _logger - Log opened: Mon May 30 19:56:12 2022 UTC\n",
      "INFO - 2022-05-30 12:56:12,145 - topics - topicmanager initialized\n"
     ]
    }
   ],
   "source": [
    "# Helper functions\n",
    "from pathlib import Path\n",
    "from copy import deepcopy\n",
    "from typing import Tuple, List, Dict\n",
    "\n",
    "import numpy as np\n",
    "import open3d as o3d\n",
    "from transforms3d.quaternions import quat2mat, mat2quat\n",
    "\n",
    "import rosbag\n",
    "from cv_bridge import CvBridge\n",
    "from geometry_msgs.msg import Transform\n",
    "\n",
    "def transform_msg2mat(transform: Transform) -> np.ndarray:\n",
    "    \"\"\"Convert a geometry_msgs/Transform to a 4x4 transformation matrix\n",
    "    :param transform: a geometry_msgs.msg.Transform.\n",
    "    \"\"\"\n",
    "    p = transform.translation\n",
    "    q = transform.rotation\n",
    "    \n",
    "    T = np.eye(4)\n",
    "    T[:3, :3] = quat2mat([q.w, q.x, q.y, q.z])  # [w, x, y, z]\n",
    "    T[:3, 3] = [p.x, p.y, p.z]\n",
    "    return T\n",
    "\n",
    "def read_frame_data_from_bag(bag: rosbag.Bag) -> Dict[str, list]:\n",
    "    \"\"\"Read frame data from a rosbag\"\"\"\n",
    "    frame_data = {\n",
    "        'rgb_imgs': [],  # [H, W, 3] uint8 np.ndarray\n",
    "        'rgb_img_timestamps': [],\n",
    "        'rgb_intrinsics': [],  # [3, 3] float64 np.ndarray\n",
    "        'rgb_intr_timestamps': [],\n",
    "        'depth_imgs': [],  # [H, W] uint16 np.ndarray\n",
    "        'depth_img_timestamps': [],\n",
    "        'depth_intrinsics': [],  # [3, 3] float64 np.ndarray\n",
    "        'depth_intr_timestamps': [],\n",
    "        'poses': [],  # T_world_cam, [4, 4] float64 np.ndarray\n",
    "        'pose_timestamps': [],\n",
    "        'timestamps': [],  # unique timestamps with imgs and poses\n",
    "    }\n",
    "    \n",
    "    cvbridge = CvBridge()\n",
    "    topic_msgs = bag.read_messages(topics=['/tf',\n",
    "                                           '/camera/color/image_raw',\n",
    "                                           '/camera/color/camera_info',\n",
    "                                           '/camera/aligned_depth_to_color/image_raw',\n",
    "                                           '/camera/aligned_depth_to_color/camera_info'])\n",
    "    for i, (topic, msg, t) in enumerate(topic_msgs):\n",
    "        # Extract msg timestamp\n",
    "        if hasattr(msg, 'header'):\n",
    "            msg_timestamp = round(msg.header.stamp.to_sec(), 6)\n",
    "        else:  # topic == '/tf'\n",
    "            msg_timestamp = round(msg.transforms[0].header.stamp.to_sec(), 6)\n",
    "\n",
    "        if topic == '/tf' and msg.transforms[0].header.frame_id == \"map\" \\\n",
    "                and msg.transforms[0].child_frame_id == 'imu':\n",
    "            frame_data['poses'].append(transform_msg2mat(msg.transforms[0].transform))\n",
    "            frame_data['pose_timestamps'].append(msg_timestamp)\n",
    "        elif topic == '/camera/color/image_raw':\n",
    "            frame_data['rgb_imgs'].append(cvbridge.imgmsg_to_cv2(msg, desired_encoding='rgb8'))\n",
    "            frame_data['rgb_img_timestamps'].append(msg_timestamp)\n",
    "        elif topic == '/camera/color/camera_info':\n",
    "            frame_data['rgb_intrinsics'].append(np.asanyarray(msg.P).reshape(3, 4)[:3, :3])\n",
    "            frame_data['rgb_intr_timestamps'].append(msg_timestamp)\n",
    "        elif topic == '/camera/aligned_depth_to_color/image_raw':\n",
    "            frame_data['depth_imgs'].append(cvbridge.imgmsg_to_cv2(msg, desired_encoding='16UC1'))\n",
    "            frame_data['depth_img_timestamps'].append(msg_timestamp)\n",
    "        elif topic == '/camera/aligned_depth_to_color/camera_info':\n",
    "            frame_data['depth_intrinsics'].append(np.asanyarray(msg.P).reshape(3, 4)[:3, :3])\n",
    "            frame_data['depth_intr_timestamps'].append(msg_timestamp)\n",
    "\n",
    "    # Find timestamps associations\n",
    "    from functools import reduce\n",
    "    timestamps = reduce(np.intersect1d, [v for k, v in frame_data.items() if '_timestamps' in k]).tolist()\n",
    "    frame_data['timestamps'] = timestamps\n",
    "    print(f\"Found {len(timestamps)} timestamps with images and poses\")\n",
    "    \n",
    "    return frame_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6f558c26",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 1488 timestamps with images and poses\n"
     ]
    }
   ],
   "source": [
    "# dense reconstruction\n",
    "bagpath = Path('/maya-slam/slam_algorithms/rosbags/d435i_bag_2022-05-20-22-23-07_posed.bag').resolve()\n",
    "bag = rosbag.Bag(bagpath)\n",
    "\n",
    "# Read frame data from rosbag\n",
    "frame_data = read_frame_data_from_bag(bag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "49eb5086",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['rgb_imgs', 'rgb_img_timestamps', 'rgb_intrinsics', 'rgb_intr_timestamps', 'depth_imgs', 'depth_img_timestamps', 'depth_intrinsics', 'depth_intr_timestamps', 'poses', 'pose_timestamps', 'timestamps'])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "frame_data.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8f4a147e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import open3d.core as o3c\n",
    "\n",
    "# Reference:\n",
    "# http://www.open3d.org/docs/0.15.1/tutorial/t_reconstruction_system/voxel_block_grid.html\n",
    "\n",
    "#device = o3c.Device('CUDA:0')\n",
    "device = o3c.Device('CPU:0')\n",
    "\n",
    "vbg = o3d.t.geometry.VoxelBlockGrid(\n",
    "    attr_names=('tsdf', 'weight', 'color'),\n",
    "    attr_dtypes=(o3c.float32, o3c.float32, o3c.float32),\n",
    "    attr_channels=((1), (1), (3)),\n",
    "    #voxel_size=3.0 / 512,  # ~5.86mm\n",
    "    voxel_size=3.0 / 300,  # ~10mm\n",
    "    block_resolution=16,\n",
    "    block_count=50000,\n",
    "    device=device\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8085c2af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<open3d.cuda.pybind.t.geometry.VoxelBlockGrid at 0x7f28b4506af0>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "vbg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b4e63e02",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_frame_data_at_timestamp(frame_data: Dict[str, list], timestamp: float):\n",
    "    \"\"\"Get frame data at timestamp\"\"\"\n",
    "    rgb_im = frame_data['rgb_imgs'][frame_data['rgb_img_timestamps'].index(timestamp)]\n",
    "    rgb_K = frame_data['rgb_intrinsics'][frame_data['rgb_intr_timestamps'].index(timestamp)]\n",
    "    \n",
    "    depth_im = frame_data['depth_imgs'][frame_data['depth_img_timestamps'].index(timestamp)]\n",
    "    depth_K = frame_data['depth_intrinsics'][frame_data['depth_intr_timestamps'].index(timestamp)]\n",
    "    \n",
    "    np.testing.assert_allclose(rgb_K, depth_K,\n",
    "                               err_msg='Different intrinsic for rgb and depth, possibly misaligned')\n",
    "    K = rgb_K\n",
    "    \n",
    "    cam_pose = frame_data['poses'][frame_data['pose_timestamps'].index(timestamp)]\\\n",
    "    \n",
    "    return rgb_im, depth_im, K, cam_pose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9ef6beb5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████| 1488/1488 [03:38<00:00,  6.82it/s]\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "for i, timestamp in enumerate(tqdm(frame_data['timestamps'])):\n",
    "    if i <= 100:\n",
    "        continue\n",
    "    \n",
    "    rgb_im, depth_im, K, cam_pose = get_frame_data_at_timestamp(frame_data, timestamp)\n",
    "    \n",
    "    rgb_im = o3d.t.geometry.Image(rgb_im).to(device)\n",
    "    depth_im = o3d.t.geometry.Image(depth_im).to(device)\n",
    "    K = o3c.Tensor(K)\n",
    "    extrinsic = o3c.Tensor(np.linalg.inv(cam_pose))\n",
    "    \n",
    "    frustum_block_coords = vbg.compute_unique_block_coordinates(\n",
    "        depth_im, K, extrinsic, depth_scale=1e3, depth_max=10.0\n",
    "    )\n",
    "    \n",
    "    vbg.integrate(frustum_block_coords, depth_im, rgb_im,\n",
    "                  K, K, extrinsic, depth_scale=1e3, depth_max=10.0)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "209090a2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PointCloud on CPU:0 [7028853 points (Float32)] Attributes: colors (dtype = Float32, shape = {7028853, 3}), normals (dtype = Float32, shape = {7028853, 3})."
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pcd = vbg.extract_point_cloud()\n",
    "pcd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "203b5833",
   "metadata": {},
   "outputs": [],
   "source": [
    "o3d.visualization.draw_geometries([pcd.to_legacy()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e53eac1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77726510",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Kernel crash\n",
    "mesh = vbg.extract_triangle_mesh()\n",
    "mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a7593cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "o3d.visualization.draw_geometries([mesh.to_legacy()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0b83d131",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'mesh' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Input \u001b[0;32mIn [10]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m o3d\u001b[38;5;241m.\u001b[39mio\u001b[38;5;241m.\u001b[39mwrite_triangle_mesh(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtest.ply\u001b[39m\u001b[38;5;124m'\u001b[39m, \u001b[43mmesh\u001b[49m\u001b[38;5;241m.\u001b[39mcuda())\n",
      "\u001b[0;31mNameError\u001b[0m: name 'mesh' is not defined"
     ]
    }
   ],
   "source": [
    "o3d.io.write_triangle_mesh('test.ply', mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8ab70fc3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<module 'open3d.cuda.pybind.io'>"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o3d.io"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "316dbff1",
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "module 'open3d' has no attribute 'cpu'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mo3d\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcpu\u001b[49m\u001b[38;5;241m.\u001b[39mpybind\u001b[38;5;241m.\u001b[39mio\u001b[38;5;241m.\u001b[39mwrite_triangle_mesh(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtest.ply\u001b[39m\u001b[38;5;124m'\u001b[39m, mesh)\n",
      "\u001b[0;31mAttributeError\u001b[0m: module 'open3d' has no attribute 'cpu'"
     ]
    }
   ],
   "source": [
    "o3d.cpu.pybind.io.write_triangle_mesh('test.ply', mesh)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7be46010",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "0f4a0897",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Image[size={480,640}, channels=1, UInt16, CUDA:0]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o3d.t.geometry.Image(depth_im).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "71bfef74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[1.0 0.0 0.0 0.0],\n",
       " [0.0 1.0 0.0 0.0],\n",
       " [0.0 0.0 1.0 0.0],\n",
       " [0.0 0.0 0.0 1.0]]\n",
       "Tensor[shape={4, 4}, stride={4, 1}, Float64, CUDA:0, 0x7f65b69e8400]"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o3c.Tensor(np.eye(4)).to(o3c.Device('CUDA:0'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "56a4784c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function open3d.cuda.pybind.core.PyCapsule.det>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "o3c.det"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2adc49d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}