mbnsf_egomotion.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d47049cb-0554-4aa7-8de4-881c788139a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install plotly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0c85c8e-8bb1-45c8-b15e-26a224fb13f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "dataset_dir = \"/mnt/local/argoverse/av1_traj_w_pose_egomotion\"\n",
    "results_dir = \"/mnt/local/argoverse/av1_traj_w_pose_egomotion/_av1_traj_w_pose_egomotion_pred\"\n",
    "log_id = \"033669d3-3d6b-3d3d-bd93-7985d86653ea\"\n",
    "\n",
    "data = np.load(f\"{dataset_dir}/{log_id}.npz\", allow_pickle=True)\n",
    "results = np.load(f\"{results_dir}/{log_id}_flow.npz\", allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fda186ea-025e-4b3b-9ff2-7c1c5b8b3d4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import plotly.graph_objects as go\n",
    "import numpy as np\n",
    "\n",
    "def plot_point_cloud(points, color, size=1.8, **kwargs):\n",
    "  return go.Scatter3d(\n",
    "      x=points[:, 0],\n",
    "      y=points[:, 1],\n",
    "      z=points[:, 2],\n",
    "      mode='markers',\n",
    "      marker=dict(size=size, color=color),\n",
    "      **kwargs,\n",
    "  )\n",
    "\n",
    "def _stack_with_none(a, b):\n",
    "  ab = np.zeros((a.shape[0], 3))\n",
    "  ab[:, 0] = a\n",
    "  ab[:, 1] = b\n",
    "  ab[:, 2] = None\n",
    "  return ab.reshape((-1,))\n",
    "\n",
    "\n",
    "def plot_flow_segments(a, b, name, **kwargs):\n",
    "  return go.Scatter3d(\n",
    "      x=_stack_with_none(a[:, 0], b[:, 0]),\n",
    "      y=_stack_with_none(a[:, 1], b[:, 1]),\n",
    "      z=_stack_with_none(a[:, 2], b[:, 2]),\n",
    "      name=name,\n",
    "      mode='lines',\n",
    "      line=kwargs,\n",
    "  )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c0b46fa-da47-4736-b3ce-88be995b1671",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Original data and inference results as is (no ego-motion compensation)\n",
    "i = 0\n",
    "p0_t0_ego = data['pcs'][i]\n",
    "f0_t0_gt_ego = data['flows'][i]\n",
    "f0_t0_pr_ego = results['flows'][i]\n",
    "p1_t1_ego = data['pcs'][i+1]\n",
    "p0_t1_gt_ego = p0_t0_ego + f0_t0_gt_ego\n",
    "p0_t1_pr_ego = p0_t0_ego + f0_t0_pr_ego\n",
    "\n",
    "fig = go.Figure(\n",
    "    data=(\n",
    "        plot_point_cloud(p0_t0_ego, color=\"red\", size=1.3, name=\"#0\"),\n",
    "        plot_point_cloud(p1_t1_ego, color=\"blue\", size=1.3, name=\"#1\"),\n",
    "        plot_flow_segments(p0_t0_ego, p0_t1_gt_ego, color=\"green\", name=\"#0 -> #1 (GT)\"),\n",
    "        plot_flow_segments(p0_t0_ego, p0_t1_pr_ego, color=\"cyan\", name=\"#0 -> #1 (PR)\")\n",
    "    ),\n",
    "    layout=dict(\n",
    "        width=1000,\n",
    "        height=1000,\n",
    "        scene=dict(aspectmode='data'),\n",
    "    )\n",
    ")\n",
    "fig.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ffe1d45-39fd-42dc-b21e-baac5f2cde87",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Post-processing to remove ego-motion\n",
    " \n",
    "# Same as SE3.transform_point_cloud in https://github.com/argoverse/argoverse-api/blob/f886ac54fba9f06f8a7d109eb663c7f501b3aa8e/argoverse/utils/se3.py#L26C5-L36C64\n",
    "def transform_point_cloud(city_from_ego: np.ndarray, point_cloud: np.ndarray) -> np.ndarray:\n",
    "    rotation = city_from_ego[:3, :3]\n",
    "    translation = city_from_ego[:3, 3]\n",
    "    return point_cloud @ rotation.T + translation\n",
    "\n",
    "city_from_ego_t0 = data['poses'][i]\n",
    "city_from_ego_t1 = data['poses'][i+1]\n",
    "p0_t0_city = transform_point_cloud(city_from_ego_t0, p0_t0_ego)\n",
    "p1_t1_city = transform_point_cloud(city_from_ego_t1, p1_t1_ego)\n",
    "# MBNSF uses ego@t1 for the flow\n",
    "p0_t1_gt_city = transform_point_cloud(city_from_ego_t1, p0_t1_gt_ego)\n",
    "p0_t1_pr_city = transform_point_cloud(city_from_ego_t1, p0_t1_pr_ego)\n",
    "\n",
    "fig = go.Figure(\n",
    "    data=(\n",
    "        plot_point_cloud(p0_t0_city, color=\"red\", size=1.3, name=\"#0\"),\n",
    "        plot_point_cloud(p1_t1_city, color=\"blue\", size=1.3, name=\"#1\"),\n",
    "        plot_flow_segments(p0_t0_city, p0_t1_gt_city, color=\"green\", name=\"#0 -> #1 (GT)\"),\n",
    "        plot_flow_segments(p0_t0_city, p0_t1_pr_city, color=\"cyan\", name=\"#0 -> #1 (PR)\")\n",
    "    ),\n",
    "    layout=dict(\n",
    "        width=1000,\n",
    "        height=1000,\n",
    "        scene=dict(aspectmode='data'),\n",
    "    )\n",
    ")\n",
    "fig.show()"
   ]
  }
 ],
 "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.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}