{"cells":[{"metadata":{},"cell_type":"markdown","source":"This notebook is a simple port of this l5kit example code\nhttps://github.com/lyft/l5kit/blob/master/examples/agent_motion_prediction/agent_motion_prediction.ipynb","execution_count":null},{"metadata":{},"cell_type":"markdown","source":"# Install and import","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install --upgrade pip\n!pip install pymap3d==2.1.0\n!pip install -U l5kit","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from typing import Dict\n\nfrom tempfile import gettempdir\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.resnet import resnet50\nfrom tqdm import tqdm\n\nfrom l5kit.configs import load_config_data\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.evaluation import write_pred_csv, compute_metrics_csv, read_gt_csv, create_chopped_dataset\nfrom l5kit.evaluation.chop_dataset import MIN_FUTURE_STEPS\nfrom l5kit.evaluation.metrics import neg_multi_log_likelihood, time_displace\nfrom l5kit.geometry import transform_points\nfrom l5kit.visualization import PREDICTED_POINTS_COLOR, TARGET_POINTS_COLOR, draw_trajectory\nfrom prettytable import PrettyTable\nfrom pathlib import Path\n\nimport os","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!cp -r ../input/lyft-motion-prediction-autonomous-vehicles ../lyft-motion-prediction-autonomous-vehicles","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prepare Data path and load cfg\nBy setting the L5KIT_DATA_FOLDER variable, we can point the script to the folder where the data lies.\n\nThen, we load our config file with relative paths and other configurations (rasteriser, training params...).","execution_count":null},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"../lyft-motion-prediction-autonomous-vehicles\"\ndm = LocalDataManager(None)\n# get config\ncfg = load_config_data(\"../input/lyft-config-files/agent_motion_config.yaml\")\nprint(cfg)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model\nOur baseline is a simple resnet50 pretrained on imagenet. We must replace the input and the final layer to address our requirements.","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load the Train Data\nOur data pipeline map a raw .zarr folder into a multi-processing instance ready for training by:\n\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","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\ntrain_cfg = cfg[\"train_data_loader\"]\nrasterizer = build_rasterizer(cfg, dm)\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset, shuffle=train_cfg[\"shuffle\"], batch_size=train_cfg[\"batch_size\"], \n                             num_workers=train_cfg[\"num_workers\"])\nprint(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== INIT MODEL\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nmodel = build_model(cfg).to(device)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\ncriterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training\nnote: if you're on MacOS and using py_satellite rasterizer, you may need to disable opencv multiprocessing by adding: 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","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== TRAIN LOOP\ntr_it = iter(train_dataloader)\nprogress_bar = tqdm(range(cfg[\"train_params\"][\"max_num_steps\"]))\nlosses_train = []\nfor _ 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)}\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Evaluation\nEvaluation 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\nIf 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\nAs such, the private test set for the competition has been \"chopped\" using the chop_dataset function.\n\n## DISCLAIMER\nWe're updating the dataset to support traffic lights. The code below has been designed to work with TLs, and it does not suppport the old interface. We expect the dataset to be online in the next few days ( disclaimer added on 08/18/20)","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== GENERATE AND LOAD CHOPPED DATASET\nnum_frames_to_chop = 100\neval_cfg = cfg[\"val_data_loader\"]\neval_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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","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\nIn this example, we simulate this pipeline by running chop_dataset on the validation set. The function stores:\n\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;","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"eval_zarr_path = str(Path(eval_base_path) / Path(dm.require(eval_cfg[\"key\"])).name)\neval_mask_path = str(Path(eval_base_path) / \"mask.npz\")\neval_gt_path = str(Path(eval_base_path) / \"gt.csv\")\n\neval_zarr = ChunkedDataset(eval_zarr_path).open()\neval_mask = np.load(eval_mask_path)[\"arr_0\"]\n# ===== INIT DATASET AND LOAD MASK\neval_dataset = AgentDataset(cfg, eval_zarr, rasterizer, agents_mask=eval_mask)\neval_dataloader = DataLoader(eval_dataset, shuffle=eval_cfg[\"shuffle\"], batch_size=eval_cfg[\"batch_size\"], \n                             num_workers=eval_cfg[\"num_workers\"])\nprint(eval_dataset)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Please note how Num Frames==(Num Scenes)*num_frames_to_chop.\n\nThe remaining frames in the scene have been sucessfully chopped off from the data","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"# ==== EVAL LOOP\nmodel.eval()\ntorch.set_grad_enabled(False)\n\n# store information for evaluation\nfuture_coords_offsets_pd = []\ntimestamps = []\n\nagent_ids = []\nprogress_bar = tqdm(eval_dataloader)\nfor data in progress_bar:\n    _, ouputs = forward(data, model, device, criterion)\n    future_coords_offsets_pd.append(ouputs.cpu().numpy().copy())\n    timestamps.append(data[\"timestamp\"].numpy().copy())\n    agent_ids.append(data[\"track_id\"].numpy().copy())","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Save results\nAfter the model has predicted trajectories for our evaluation set, we can save them in a csv file.\n\nDuring the competition, only the .zarr and the mask will be provided for the private test set evaluation. Your solution is expected to generate a csv file which will be compared to the ground truth one on a separate server","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"pred_path = f\"{gettempdir()}/pred.csv\"\n\nwrite_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              )","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Perform Evaluation\nPleae 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\nIn 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","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"metrics = compute_metrics_csv(eval_gt_path, pred_path, [neg_multi_log_likelihood, time_displace])\nfor metric_name, metric_mean in metrics.items():\n    print(metric_name, metric_mean)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## Visualise Results\nWe can also visualise some results from the ego (AV) point of view for those frames of interest (the 100th of each scene).\n\nHowever, 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","execution_count":null},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\ntorch.set_grad_enabled(False)\n\n# build a dict to retrieve future trajectories from GT\ngt_rows = {}\nfor row in read_gt_csv(eval_gt_path):\n    gt_rows[row[\"track_id\"] + row[\"timestamp\"]] = row[\"coord\"]\n\neval_ego_dataset = EgoDataset(cfg, eval_dataset.dataset, rasterizer)\n\nfor 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(out_pos + data_agent[\"centroid\"][:2])\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[\"world_to_image\"])\n    target_positions = transform_points(np.concatenate(target_positions), data_ego[\"world_to_image\"])\n\n    yaws = np.zeros((len(predicted_positions), 1))\n    draw_trajectory(im_ego, predicted_positions, yaws, PREDICTED_POINTS_COLOR)\n    draw_trajectory(im_ego, target_positions, yaws, TARGET_POINTS_COLOR)\n\n    plt.imshow(im_ego[::-1])\n    plt.show()","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}