{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c888728a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58a759f3",
   "metadata": {},
   "source": [
    "# Devkit\n",
    "\n",
    "Please refer to [the OpenLane-V2 devkit](https://github.com/OpenDriveLab/OpenLane-V2/blob/master/docs/devkit.md) for a detailed description of the API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "25832f40",
   "metadata": {},
   "outputs": [],
   "source": [
    "root_path = './data/OpenLane-V2'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efe74073",
   "metadata": {},
   "source": [
    "## Preprocessing\n",
    "\n",
    "Collect meta data of a collection of data and store it into a pickle file, where the collection can be user-defined by modifying the input data_dict. Note that for evaluation, we subsample the number of points of lane centerlines with an interval of 20. For training, the `point_interval` can be set to 1 to disable subsampling.\n",
    "\n",
    "To preprocess the whole dataset, please refer to [preprocess](https://github.com/OpenDriveLab/OpenLane-V2/blob/master/data/README.md#preprocess)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "ccc41f71",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "collecting data_dict_sample: 100%|██████████████| 64/64 [00:01<00:00, 43.61it/s]\n"
     ]
    }
   ],
   "source": [
    "from openlanev2.io import io\n",
    "from openlanev2.preprocessing import collect\n",
    "\n",
    "data_dict = io.json_load(f'{root_path}/data_dict_sample.json')\n",
    "collect(root_path, data_dict, 'data_dict_sample', point_interval=20)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "667aae21",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Dataset\n",
    "\n",
    "Load meta data from the generated pickle file.\n",
    "A frame can be fetched by specifying the identifier `(split, segment_id, timestamp)` or giving the index of the data in the collection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4a41f652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openlanev2.dataset import Collection, Frame\n",
    "\n",
    "collection = Collection(root_path, root_path, 'data_dict_sample')\n",
    "frame = collection.get_frame_via_identifier(('train', '00492', '315970276749927222'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41f764e9",
   "metadata": {},
   "source": [
    "The pose of the ego vehicle can be obtained by `get_pose()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "82a18679",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rotation \n",
      " [[-0.03161852  0.99932457  0.01872611]\n",
      " [-0.99904989 -0.03103652 -0.0305949 ]\n",
      " [-0.02999304 -0.01967568  0.99935644]] \n",
      "\n",
      "translation \n",
      " [6657.19639943 1842.86566032   59.89421563] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "for k, v in frame.get_pose().items():\n",
    "    print(k, '\\n', v, '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "251a0b51",
   "metadata": {},
   "source": [
    "Each frame contains multi-view images covering the whole horizontal FOV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ea7df923",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ring_front_center\n",
      "ring_front_left\n",
      "ring_front_right\n",
      "ring_rear_left\n",
      "ring_rear_right\n",
      "ring_side_left\n",
      "ring_side_right\n"
     ]
    }
   ],
   "source": [
    "for k in frame.get_camera_list():\n",
    "    print(k)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6b621d6",
   "metadata": {},
   "source": [
    "For each image, we provide intrinsic and extrinsic matrices of the camera."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b6bf5537",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 632880 characters]\n",
      "text/plain": [
       "<Figure size 1600x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intrinsic - K \n",
      " [[1.77753967e+03 0.00000000e+00 7.77762878e+02]\n",
      " [0.00000000e+00 1.77753967e+03 1.01631311e+03]\n",
      " [0.00000000e+00 0.00000000e+00 1.00000000e+00]] \n",
      "\n",
      "intrinsic - distortion \n",
      " [-0.24479243 -0.19577468  0.30131747] \n",
      "\n",
      "extrinsic - rotation \n",
      " [[-8.48589870e-04  1.00005773e-02  9.99949633e-01]\n",
      " [-9.99998983e-01 -1.15429954e-03 -8.37087508e-04]\n",
      " [ 1.14587004e-03 -9.99949327e-01  1.00015467e-02]] \n",
      "\n",
      "extrinsic - translation \n",
      " [1.63315125 0.00800013 1.38385219] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "camera = frame.get_camera_list()[0]\n",
    "\n",
    "image = frame.get_rgb_image(camera)\n",
    "\n",
    "plt.figure(figsize=[16, 9])\n",
    "plt.imshow(image)\n",
    "plt.show()\n",
    "\n",
    "meta = {\n",
    "    'intrinsic': frame.get_intrinsic(camera),\n",
    "    'extrinsic': frame.get_extrinsic(camera),\n",
    "}\n",
    "for key, value in meta.items():\n",
    "    for k, v in value.items():\n",
    "        print(key, '-', k, '\\n', v, '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3110001f",
   "metadata": {},
   "source": [
    "Annotations of a frame can be loaded separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b6831a32",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['lane_centerline', 'traffic_element', 'topology_lclc', 'topology_lcte']\n",
      "\n",
      "lane_centerline:\n",
      "id <class 'str'>\n",
      "points (11, 3)\n",
      "is_intersection_or_connector <class 'bool'>\n",
      "\n",
      "traffic_element:\n",
      "id <class 'str'>\n",
      "category <class 'int'>\n",
      "attribute <class 'int'>\n",
      "points (2, 2)\n",
      "\n",
      "topology_lclc:\n",
      "adjacent matrix (61, 61)\n",
      "\n",
      "topology_lcte:\n",
      "adjacent matrix (61, 11)\n"
     ]
    }
   ],
   "source": [
    "annotations = frame.get_annotations()\n",
    "print(list(annotations.keys()))\n",
    "\n",
    "for a, anno in zip(\n",
    "    annotations.keys(), \n",
    "    [\n",
    "        frame.get_annotations_lane_centerlines()[0],\n",
    "        frame.get_annotations_traffic_elements()[0],\n",
    "        frame.get_annotations_topology_lclc(),\n",
    "        frame.get_annotations_topology_lcte(),\n",
    "    ]):\n",
    "    print(f'\\n{a}:')\n",
    "    if isinstance(anno, dict):\n",
    "        for k, v in anno.items():\n",
    "            print(k, v.shape) if isinstance(v, np.ndarray) else print(k, type(v))\n",
    "    else:\n",
    "        print('adjacent matrix', anno.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eafe8f3f",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Visualization\n",
    "\n",
    "Preprocess the annotations for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bd29fbd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openlanev2.visualization import draw_annotation_bev, draw_annotation_pv, assign_attribute, assign_topology\n",
    "\n",
    "annotations = assign_attribute(annotations)\n",
    "annotations = assign_topology(annotations)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b68db7ce",
   "metadata": {},
   "source": [
    "Visualize annotations in the bird's-eye view. `with_attribute` denotes that color is used to draw lane centerlines, which indicates the attribute of the lane centerline, or not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "57125f57",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 109648 characters]\n",
      "text/plain": [
       "<Figure size 1600x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image_bev = draw_annotation_bev(\n",
    "    annotations, \n",
    "    with_attribute=True,\n",
    ")\n",
    "\n",
    "plt.figure(figsize=[16, 9])\n",
    "plt.imshow(image_bev)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "454117c7",
   "metadata": {},
   "source": [
    "Visualize annotations in the perspective view. Note that as traffic elements are annotated in the front-view image, visualization for traffic elements and the topology is only valid in the front-view image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e071b2fd",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 667140 characters]\n",
      "text/plain": [
       "<Figure size 1600x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image_pv = draw_annotation_pv(\n",
    "    camera, \n",
    "    image.copy(), \n",
    "    annotations,\n",
    "    meta['intrinsic'],\n",
    "    meta['extrinsic'],\n",
    "    with_attribute=True, \n",
    "    with_topology=True,\n",
    ")\n",
    "\n",
    "plt.figure(figsize=[16, 9])\n",
    "plt.imshow(image_pv)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9acb0ebd",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## Evaluation\n",
    "\n",
    "We provide tools for evaluation based on the generated pickle files.\n",
    "The input can be a path to the pickle or a formatted dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f08343c9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Dummy evaluation on ground truth.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "calculating distances:: 100%|███████████████████| 64/64 [00:17<00:00,  3.69it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "OpenLane-V2 Score - 1.0\n",
      "    DET_l - 1.0\n",
      "    DET_t - 1.0\n",
      "    TOP_ll - 1.0\n",
      "    TOP_lt - 1.0\n",
      "F-Score for 3D Lane - 1.0\n"
     ]
    }
   ],
   "source": [
    "from openlanev2.evaluation import evaluate\n",
    "from openlanev2.utils import format_metric\n",
    "\n",
    "format_metric(evaluate(ground_truth=f'{root_path}/data_dict_sample.pkl', predictions=None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "479e6c30",
   "metadata": {},
   "source": [
    "To construct results for evaluation, please follow the example below.\n",
    "Some keys are necessary for a valid evaluation.\n",
    "Note that `openlanev2.preprocessing.check_results` conducting checking on prediction format in the `evaluate` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3c04ca67",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "*** Missing key [country / region] for a valid submission ***\n",
      "\n",
      "\n",
      "*** Missing key [authors] for a valid submission ***\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "calculating distances:: 100%|███████████████████| 64/64 [00:01<00:00, 59.32it/s]\n"
     ]
    }
   ],
   "source": [
    "results = {\n",
    "    'method': 'dummy',\n",
    "    'e-mail': 'dummy@openlanev2',\n",
    "    'institution / company': 'dummy',\n",
    "    'results': {},\n",
    "}\n",
    "for key in collection.keys:\n",
    "    results['results'][key] = {\n",
    "        'predictions': {\n",
    "            'lane_centerline': [{\n",
    "                'id': 20,\n",
    "                'points': np.asarray([\n",
    "                    [1, 2, 3],\n",
    "                    [2, 3, 4],\n",
    "                ]),\n",
    "                'confidence': 0.6,\n",
    "            }],\n",
    "            'traffic_element': [{\n",
    "                'id': i,\n",
    "                'category': 0,\n",
    "                'attribute': i,\n",
    "                'points': np.asarray([\n",
    "                    [1, 2],\n",
    "                    [3, 4],\n",
    "                ]),\n",
    "                'confidence': 0.7,\n",
    "            } for i in range(13)],\n",
    "            'topology_lclc': np.asarray([\n",
    "                [0],\n",
    "            ]),\n",
    "            'topology_lcte': np.asarray([\n",
    "                [0.8 for _ in range(13)],\n",
    "            ]),\n",
    "        }\n",
    "    }\n",
    "\n",
    "results = evaluate(\n",
    "    ground_truth=f'{root_path}/data_dict_sample.pkl', \n",
    "    predictions=results,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65b7766d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "# Train a Model\n",
    "\n",
    "To train your first model on OpenLane-V2, please checkout to [plug-ins](https://github.com/OpenDriveLab/OpenLane-V2#train-a-model) for the provided baseline model in prevailing training frameworks. We currently provide a plug-in to [mmdet3d](https://github.com/open-mmlab/mmdetection3d), in which a pipeline for training, evaluation, visualization, and formatting submission files is built.\n",
    "\n",
    "Your joint contribution to maintaining or adding new plug-ins to other frameworks is welcomed."
   ]
  }
 ],
 "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.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
