{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom typing import Dict\nfrom tempfile import gettempdir\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn \nfrom torch.utils.data import DataLoader\nimport torchvision.models as models\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(device)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"import l5kit\nprint(l5kit.__version__)\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\n\nfrom l5kit.evaluation import write_pred_csv, compute_metrics_csv, read_gt_csv, create_chopped_dataset, write_gt_csv\nfrom l5kit.evaluation.chop_dataset import MIN_FUTURE_STEPS\nfrom l5kit.evaluation.metrics import neg_multi_log_likelihood, time_displace\n\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\nos.environ[\"L5KIT_DATA_FOLDER\"] = \"../input/lyft-motion-prediction-autonomous-vehicles\"\ndm = LocalDataManager(None)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {'model_architecture': 'resnet18',\n                     'history_num_frames': 10,\n                     'history_step_size': 1,\n                     'history_delta_time': 0.1,\n                     'future_num_frames': 50,\n                     'future_step_size': 1,\n                     'future_delta_time': 0.1},\n    'raster_params': {'raster_size': [224, 224],\n                      'pixel_size': [0.5, 0.5],\n                      'ego_center': [0.25, 0.5],\n                      'map_type': 'py_semantic',\n                      'satellite_map_key': 'aerial_map/aerial_map.png',\n                      'semantic_map_key': 'semantic_map/semantic_map.pb',\n                      'dataset_meta_key': 'meta.json',\n                      'filter_agents_threshold': 0.5,\n                     'disable_traffic_light_faces':False},\n    'example_data_loader': {'key': 'scenes/sample.zarr',\n                        'batch_size': 4,\n                        'shuffle': False,\n                        'num_workers': 0}}\n\nrasterizer = build_rasterizer(cfg, dm)\n\nexample_cfg = cfg[\"example_data_loader\"]\nexample_zarr = ChunkedDataset(dm.require(example_cfg[\"key\"])).open()\ndataset1 = AgentDataset(cfg, example_zarr, rasterizer)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## load model"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftMultiModel(nn.Module):\n\n    def __init__(self, cfg: Dict, num_modes=3):\n        super().__init__()\n        \n        \"\"\"backbone network\"\"\"\n        self.backbone = models.resnet18()\n        self.backbone_out_features = 512  #for resnet18 and resnet34; 2048 for the other resnets\n        \n        \n        self.num_history_channels = (cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        self.num_in_channels = 3 + self.num_history_channels\n        \n        self.num_modes = num_modes\n        self.num_targets = cfg[\"model_params\"][\"future_num_frames\"]\n        self.num_preds = 2 * self.num_targets * self.num_modes \n        \n\n        self.backbone.conv1 = nn.Conv2d(\n            self.num_in_channels,\n            self.backbone.conv1.out_channels,\n            kernel_size=self.backbone.conv1.kernel_size,\n            stride=self.backbone.conv1.stride,\n            padding=self.backbone.conv1.padding,\n            bias=False,\n        )\n\n        self.head = nn.Sequential(\n            nn.Linear(in_features=self.backbone_out_features, out_features=4096),\n            nn.Linear(4096, out_features = self.num_preds + self.num_modes)\n        )\n\n        \n    def forward(self, x):\n        x = self.backbone.conv1(x)\n        x = self.backbone.bn1(x)\n        x = self.backbone.relu(x)\n        x = self.backbone.maxpool(x)\n\n        x = self.backbone.layer1(x)\n        x = self.backbone.layer2(x)\n        x = self.backbone.layer3(x)\n        x = self.backbone.layer4(x)\n\n        x = self.backbone.avgpool(x)\n        x = torch.flatten(x, 1)\n\n        x = self.head(x)\n\n        confidences, preds = torch.split(x, [self.num_modes, self.num_preds], dim=1)\n      \n        return  confidences, preds.view(-1, self.num_modes , self.num_targets, 2) \n    \nmodel = LyftMultiModel(cfg).to(device)\nmodel.load_state_dict(torch.load('../input/lyft-motion-resnet18-multi-mode-training-5/lyft_multi_mode_model5_cpu.pth'))\nmodel.eval()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## plot predictions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_target(dataset, idx, COLOR = TARGET_POINTS_COLOR):\n    data = dataset[idx]\n\n    # convert raster into rgb\n    im = dataset.rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0)) \n    # transform points from agent local coordinate to world coordinate\n    #positions_in_world = transform_points(data[\"target_positions\"], data[\"world_from_agent\"])\n    # transform from world coordinates onto raster\n    #positions_in_raster = transform_points(positions_in_world, data[\"raster_from_world\"])\n    positions_in_raster = transform_points(data[\"target_positions\"], data[\"raster_from_agent\"])\n    draw_trajectory(im, positions_in_raster, COLOR, yaws=data[\"target_yaws\"])\n    plt.imshow(im[::-1])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def plot_predictions(dataset, i, model):\n    # visualize predictions\n    data = dataset[i]\n    \n    x = torch.FloatTensor(data[\"image\"]).unsqueeze(0).to(device)\n    confidences, preds = model(x)\n    conf_probs = torch.softmax(confidences, 1)\n\n    conf_probs = conf_probs.squeeze()\n    preds = preds.squeeze()\n\n    fig = plt.figure( figsize=(16, 16))\n\n    ax = fig.add_subplot(131)\n    idx = 0\n    ax.set_title(str(round(conf_probs[idx].item(), 4)))\n\n    im = dataset.rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0))  # convert raster into rgb\n    # transform points from agent local coordinate onto raster\n    positions_in_raster = transform_points(preds[idx].cpu().detach().numpy(), data[\"raster_from_agent\"])\n    draw_trajectory(im, positions_in_raster, PREDICTED_POINTS_COLOR, yaws=data[\"target_yaws\"])\n    ax.imshow(im[::-1])\n\n\n    ax = fig.add_subplot(132)\n    idx = 1\n    ax.set_title(str(round(conf_probs[idx].item(), 4)))\n\n    im = dataset.rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0))  # convert raster into rgb\n    # transform points from agent local coordinate onto raster\n    positions_in_raster = transform_points(preds[idx].cpu().detach().numpy(), data[\"raster_from_agent\"])\n    draw_trajectory(im, positions_in_raster, (255, 0 , 0), yaws=data[\"target_yaws\"])\n    ax.imshow(im[::-1])\n\n\n    ax = fig.add_subplot(133)\n    idx = 2\n\n    ax.set_title(str(round(conf_probs[idx].item(), 4)))\n    im = dataset.rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0))  # convert raster into rgb\n    # transform points from agent local coordinate onto raster\n    positions_in_raster = transform_points(preds[idx].cpu().detach().numpy(), data[\"raster_from_agent\"])\n    draw_trajectory(im, positions_in_raster, (255, 255 , 0), yaws=data[\"target_yaws\"])\n    ax.imshow(im[::-1])\n\n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"idx = 500\nplot_target(dataset1, idx)\nplot_predictions(dataset1,idx , model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"idx = 1000\nplot_target(dataset1, idx)\nplot_predictions(dataset1,idx , model)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"idx = 2500\nplot_target(dataset1, idx)\nplot_predictions(dataset1,idx , model)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## animation"},{"metadata":{"trusted":true},"cell_type":"code","source":"from matplotlib import animation, rc\nfrom IPython.display import HTML\nfrom IPython.display import display, clear_output\n\nimport PIL\n\nrc('animation', html='jshtml')\n\ndef animate_solution(images):\n\n    def animate(i):\n        im.set_data(images[i])\n \n    fig, ax = plt.subplots()\n    im = ax.imshow(images[0])\n    \n    return animation.FuncAnimation(fig, animate, frames=len(images), interval=60)\n\nego_dataset = EgoDataset(cfg, example_zarr, rasterizer)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"sence 8, 14"},{"metadata":{"trusted":true},"cell_type":"code","source":"scene_idx = 25\nindexes = ego_dataset.get_scene_indices(scene_idx)\nprint(len(indexes))\nplot_target(ego_dataset, indexes[0])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def animate_prediction(images, confs):\n\n    def animate(i):\n        im0.set_data(images[i][0])\n        im1.set_data(images[i][1])\n        im2.set_data(images[i][2])\n        im3.set_data(images[i][3])\n        axes[0].set_title(str(round(confs[i,0], 2)))\n        axes[1].set_title(str(round(confs[i,1], 2)))\n        axes[2].set_title(str(round(confs[i,2], 2)))\n        axes[3].set_title('ground truth')\n \n    fig, axes = plt.subplots(1,4, figsize=(16,4))\n\n    im0 = axes[0].imshow(images[0][0])\n    im1 = axes[1].imshow(images[0][1])\n    im2 = axes[2].imshow(images[0][2])\n    im3 = axes[3].imshow(images[0][3])\n    \n    return animation.FuncAnimation(fig, animate, frames=len(images), interval=60*4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"COLORS = [PREDICTED_POINTS_COLOR, (255, 0 , 0), (255, 255 , 0)]\n\ndef make_anim(scene_idx):\n    indexes = ego_dataset.get_scene_indices(scene_idx)\n    images, confs = [], []\n\n    for idx in indexes:\n        ims = []\n        data = ego_dataset[idx]\n    \n        x = torch.FloatTensor(data[\"image\"]).unsqueeze(0).to(device)\n        confidences, preds = model(x)\n        conf_probs = torch.softmax(confidences, 1)\n        preds = preds.squeeze()\n    \n        for idx in range(3):\n            im = ego_dataset.rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0)) \n            positions_in_raster = transform_points(preds[idx].cpu().detach().numpy(), data[\"raster_from_agent\"])\n            draw_trajectory(im, positions_in_raster, COLORS[idx], yaws=data[\"target_yaws\"])\n            ims.append(PIL.Image.fromarray(im[::-1]))\n    \n        # target\n        im = ego_dataset.rasterizer.to_rgb(data[\"image\"].transpose(1, 2, 0)) \n        positions_in_raster = transform_points(data[\"target_positions\"], data[\"raster_from_agent\"])\n        draw_trajectory(im, positions_in_raster, TARGET_POINTS_COLOR, yaws=data[\"target_yaws\"])\n        ims.append(PIL.Image.fromarray(im[::-1]))\n    \n        confs.append(conf_probs.detach().numpy())\n        images.append(ims)\n    \n    confs = np.concatenate(confs)\n    return animate_prediction(images, confs)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"anim = make_anim(8)\nHTML(anim.to_jshtml())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"anim = make_anim(14)\nHTML(anim.to_jshtml())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"anim = make_anim(25)\nHTML(anim.to_jshtml())","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}