examples/agent_motion_prediction/agent_motion_prediction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Dict\n",
    "\n",
    "from tempfile import gettempdir\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn, optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision.models.resnet import resnet50\n",
    "from tqdm import tqdm\n",
    "\n",
    "from l5kit.configs import load_config_data\n",
    "from l5kit.data import LocalDataManager, ChunkedDataset\n",
    "from l5kit.dataset import AgentDataset, EgoDataset\n",
    "from l5kit.rasterization import build_rasterizer\n",
    "from l5kit.evaluation import write_pred_csv, compute_metrics_csv, read_gt_csv, create_chopped_dataset\n",
    "from l5kit.evaluation.chop_dataset import MIN_FUTURE_STEPS\n",
    "from l5kit.evaluation.metrics import neg_multi_log_likelihood, time_displace\n",
    "from l5kit.geometry import transform_points\n",
    "from l5kit.visualization import PREDICTED_POINTS_COLOR, TARGET_POINTS_COLOR, draw_trajectory\n",
    "from prettytable import PrettyTable\n",
    "from pathlib import Path\n",
    "\n",
    "import os"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare Data path and load cfg\n",
    "\n",
    "By setting the `L5KIT_DATA_FOLDER` variable, we can point the script to the folder where the data lies.\n",
    "\n",
    "Then, we load our config file with relative paths and other configurations (rasteriser, training params...)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set env variable for data\n",
    "os.environ[\"L5KIT_DATA_FOLDER\"] = \"PATH_TO_DATA\"\n",
    "dm = LocalDataManager(None)\n",
    "# get config\n",
    "cfg = load_config_data(\"./agent_motion_config.yaml\")\n",
    "print(cfg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model\n",
    "\n",
    "Our baseline is a simple `resnet50` pretrained on `imagenet`. We must replace the input and the final layer to address our requirements."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(cfg: Dict) -> torch.nn.Module:\n",
    "    # load pre-trained Conv2D model\n",
    "    model = resnet50(pretrained=True)\n",
    "\n",
    "    # change input channels number to match the rasterizer's output\n",
    "    num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n",
    "    num_in_channels = 3 + num_history_channels\n",
    "    model.conv1 = nn.Conv2d(\n",
    "        num_in_channels,\n",
    "        model.conv1.out_channels,\n",
    "        kernel_size=model.conv1.kernel_size,\n",
    "        stride=model.conv1.stride,\n",
    "        padding=model.conv1.padding,\n",
    "        bias=False,\n",
    "    )\n",
    "    # change output size to (X, Y) * number of future states\n",
    "    num_targets = 2 * cfg[\"model_params\"][\"future_num_frames\"]\n",
    "    model.fc = nn.Linear(in_features=2048, out_features=num_targets)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(data, model, device, criterion):\n",
    "    inputs = data[\"image\"].to(device)\n",
    "    target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n",
    "    targets = data[\"target_positions\"].to(device)\n",
    "    # Forward pass\n",
    "    outputs = model(inputs).reshape(targets.shape)\n",
    "    loss = criterion(outputs, targets)\n",
    "    # not all the output steps are valid, but we can filter them out from the loss using availabilities\n",
    "    loss = loss * target_availabilities\n",
    "    loss = loss.mean()\n",
    "    return loss, outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the Train Data\n",
    "\n",
    "Our data pipeline map a raw `.zarr` folder into a multi-processing instance ready for training by:\n",
    "- loading the `zarr` into a `ChunkedDataset` object. This object has a reference to the different arrays into the zarr (e.g. agents and traffic lights);\n",
    "- wrapping the `ChunkedDataset` into an `AgentDataset`, which inherits from torch `Dataset` class;\n",
    "- passing the `AgentDataset` into a torch `DataLoader`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===== INIT DATASET\n",
    "train_cfg = cfg[\"train_data_loader\"]\n",
    "rasterizer = build_rasterizer(cfg, dm)\n",
    "train_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n",
    "train_dataset = AgentDataset(cfg, train_zarr, rasterizer)\n",
    "train_dataloader = DataLoader(train_dataset, shuffle=train_cfg[\"shuffle\"], batch_size=train_cfg[\"batch_size\"], \n",
    "                             num_workers=train_cfg[\"num_workers\"])\n",
    "print(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==== INIT MODEL\n",
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = build_model(cfg).to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\n",
    "criterion = nn.MSELoss(reduction=\"none\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training\n",
    "\n",
    "note: if you're on MacOS and using `py_satellite` rasterizer, you may need to disable opencv multiprocessing by adding:\n",
    "`cv2.setNumThreads(0)` before the following cell. This seems to only affect running in python notebook and it's caused by the `cv2.warpaffine` function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==== TRAIN LOOP\n",
    "tr_it = iter(train_dataloader)\n",
    "progress_bar = tqdm(range(cfg[\"train_params\"][\"max_num_steps\"]))\n",
    "losses_train = []\n",
    "for _ in progress_bar:\n",
    "    try:\n",
    "        data = next(tr_it)\n",
    "    except StopIteration:\n",
    "        tr_it = iter(train_dataloader)\n",
    "        data = next(tr_it)\n",
    "    model.train()\n",
    "    torch.set_grad_enabled(True)\n",
    "    loss, _ = forward(data, model, device, criterion)\n",
    "\n",
    "    # Backward pass\n",
    "    optimizer.zero_grad()\n",
    "    loss.backward()\n",
    "    optimizer.step()\n",
    "\n",
    "    losses_train.append(loss.item())\n",
    "    progress_bar.set_description(f\"loss: {loss.item()} loss(avg): {np.mean(losses_train)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plot Loss Curve\n",
    "We can plot the train loss against the iterations (batch-wise)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(np.arange(len(losses_train)), losses_train, label=\"train loss\")\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluation\n",
    "\n",
    "Evaluation follows a slightly different protocol than training. When working with time series, we must be absolutely sure to avoid leaking the future in the data.\n",
    "\n",
    "If we followed the same protocol of training, one could just read ahead in the `.zarr` and forge a perfect solution at run-time, even for a private test set.\n",
    "\n",
    "As such, **the private test set for the competition has been \"chopped\" using the `chop_dataset` function**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ===== GENERATE AND LOAD CHOPPED DATASET\n",
    "num_frames_to_chop = 100\n",
    "eval_cfg = cfg[\"val_data_loader\"]\n",
    "eval_base_path = create_chopped_dataset(dm.require(eval_cfg[\"key\"]), cfg[\"raster_params\"][\"filter_agents_threshold\"], \n",
    "                              num_frames_to_chop, cfg[\"model_params\"][\"future_num_frames\"], MIN_FUTURE_STEPS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The result is that **each scene has been reduced to only 100 frames**, and **only valid agents in the 100th frame will be used to compute the metrics**. Because following frames in the scene have been chopped off, we can't just look ahead to get the future of those agents.\n",
    "\n",
    "In this example, we simulate this pipeline by running `chop_dataset` on the validation set. The function stores:\n",
    "- a new chopped `.zarr` dataset, in which each scene has only the first 100 frames;\n",
    "- a numpy mask array where only valid agents in the 100th frame are True;\n",
    "- a ground-truth file with the future coordinates of those agents;\n",
    "\n",
    "Please note how the total number of frames is now equal to the number of scenes multipled by `num_frames_to_chop`. \n",
    "\n",
    "The remaining frames in the scene have been sucessfully chopped off from the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_zarr_path = str(Path(eval_base_path) / Path(dm.require(eval_cfg[\"key\"])).name)\n",
    "eval_mask_path = str(Path(eval_base_path) / \"mask.npz\")\n",
    "eval_gt_path = str(Path(eval_base_path) / \"gt.csv\")\n",
    "\n",
    "eval_zarr = ChunkedDataset(eval_zarr_path).open()\n",
    "eval_mask = np.load(eval_mask_path)[\"arr_0\"]\n",
    "# ===== INIT DATASET AND LOAD MASK\n",
    "eval_dataset = AgentDataset(cfg, eval_zarr, rasterizer, agents_mask=eval_mask)\n",
    "eval_dataloader = DataLoader(eval_dataset, shuffle=eval_cfg[\"shuffle\"], batch_size=eval_cfg[\"batch_size\"], \n",
    "                             num_workers=eval_cfg[\"num_workers\"])\n",
    "print(eval_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Storing Predictions\n",
    "There is a small catch to be aware of when saving the model predictions. The output of the models are coordinates in `agent` space and we need to convert them into displacements in `world` space.\n",
    "\n",
    "To do so, we first convert them back into the `world` space and we then subtract the centroid coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ==== EVAL LOOP\n",
    "model.eval()\n",
    "torch.set_grad_enabled(False)\n",
    "\n",
    "# store information for evaluation\n",
    "future_coords_offsets_pd = []\n",
    "timestamps = []\n",
    "agent_ids = []\n",
    "\n",
    "progress_bar = tqdm(eval_dataloader)\n",
    "for data in progress_bar:\n",
    "    _, ouputs = forward(data, model, device, criterion)\n",
    "    \n",
    "    # convert agent coordinates into world offsets\n",
    "    agents_coords = ouputs.cpu().numpy()\n",
    "    world_from_agents = data[\"world_from_agent\"].numpy()\n",
    "    centroids = data[\"centroid\"].numpy()\n",
    "    coords_offset = transform_points(agents_coords, world_from_agents) - centroids[:, None, :2]\n",
    "    \n",
    "    future_coords_offsets_pd.append(np.stack(coords_offset))\n",
    "    timestamps.append(data[\"timestamp\"].numpy().copy())\n",
    "    agent_ids.append(data[\"track_id\"].numpy().copy())\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Save results\n",
    "After the model has predicted trajectories for our evaluation set, we can save them in a `csv` file.\n",
    "\n",
    "During the competition, only the `.zarr` and the mask will be provided for the private test set evaluation.\n",
    "Your solution is expected to generate a csv file which will be compared to the ground truth one on a separate server"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_path = f\"{gettempdir()}/pred.csv\"\n",
    "\n",
    "write_pred_csv(pred_path,\n",
    "               timestamps=np.concatenate(timestamps),\n",
    "               track_ids=np.concatenate(agent_ids),\n",
    "               coords=np.concatenate(future_coords_offsets_pd),\n",
    "              )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Perform Evaluation\n",
    "Pleae note that our metric supports multi-modal predictions (i.e. multiple predictions for a single GT trajectory). In that case, you will need to provide a confidence for each prediction (confidences must all be between 0 and 1 and sum to 1).\n",
    "\n",
    "In this simple example we don't generate multiple trajectories, so we won't pass any confidences vector. Internally, the metric computation will assume a single trajectory with confidence equal to 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metrics = compute_metrics_csv(eval_gt_path, pred_path, [neg_multi_log_likelihood, time_displace])\n",
    "for metric_name, metric_mean in metrics.items():\n",
    "    print(metric_name, metric_mean)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualise Results\n",
    "We can also visualise some results from the ego (AV) point of view for those frames of interest (the 100th of each scene).\n",
    "\n",
    "However, as we chopped off the future from the dataset **we must use the GT csv if we want to plot the future trajectories of the agents**\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.eval()\n",
    "torch.set_grad_enabled(False)\n",
    "\n",
    "# build a dict to retrieve future trajectories from GT\n",
    "gt_rows = {}\n",
    "for row in read_gt_csv(eval_gt_path):\n",
    "    gt_rows[row[\"track_id\"] + row[\"timestamp\"]] = row[\"coord\"]\n",
    "\n",
    "eval_ego_dataset = EgoDataset(cfg, eval_dataset.dataset, rasterizer)\n",
    "\n",
    "for frame_number in range(99, len(eval_zarr.frames), 100):  # start from last frame of scene_0 and increase by 100\n",
    "    agent_indices = eval_dataset.get_frame_indices(frame_number) \n",
    "    if not len(agent_indices):\n",
    "        continue\n",
    "\n",
    "    # get AV point-of-view frame\n",
    "    data_ego = eval_ego_dataset[frame_number]\n",
    "    im_ego = rasterizer.to_rgb(data_ego[\"image\"].transpose(1, 2, 0))\n",
    "    center = np.asarray(cfg[\"raster_params\"][\"ego_center\"]) * cfg[\"raster_params\"][\"raster_size\"]\n",
    "    \n",
    "    predicted_positions = []\n",
    "    target_positions = []\n",
    "\n",
    "    for v_index in agent_indices:\n",
    "        data_agent = eval_dataset[v_index]\n",
    "\n",
    "        out_net = model(torch.from_numpy(data_agent[\"image\"]).unsqueeze(0).to(device))\n",
    "        out_pos = out_net[0].reshape(-1, 2).detach().cpu().numpy()\n",
    "        # store absolute world coordinates\n",
    "        predicted_positions.append(transform_points(out_pos, data_agent[\"world_from_agent\"]))\n",
    "        # retrieve target positions from the GT and store as absolute coordinates\n",
    "        track_id, timestamp = data_agent[\"track_id\"], data_agent[\"timestamp\"]\n",
    "        target_positions.append(gt_rows[str(track_id) + str(timestamp)] + data_agent[\"centroid\"][:2])\n",
    "\n",
    "\n",
    "    # convert coordinates to AV point-of-view so we can draw them\n",
    "    predicted_positions = transform_points(np.concatenate(predicted_positions), data_ego[\"raster_from_world\"])\n",
    "    target_positions = transform_points(np.concatenate(target_positions), data_ego[\"raster_from_world\"])\n",
    "\n",
    "    draw_trajectory(im_ego, predicted_positions, PREDICTED_POINTS_COLOR)\n",
    "    draw_trajectory(im_ego, target_positions, TARGET_POINTS_COLOR)\n",
    "\n",
    "    plt.imshow(im_ego)\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.3"
  },
  "pycharm": {
   "stem_cell": {
    "cell_type": "raw",
    "metadata": {
     "collapsed": false
    },
    "source": []
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}