{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip uninstall -y typing\n!pip install l5kit","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import numpy as np\nimport os\nimport torch\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.resnet import resnet50\nfrom tqdm import tqdm\nfrom typing import Dict\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.evaluation import write_pred_csv\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.geometry import transform_points","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n\n# Training notebook's output.\nWEIGHT_FILE = \"/kaggle/input/model-resnet50/model_resnet50.pth\"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\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    },\n    \n    'raster_params': {\n        'raster_size': [300, 300],\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    },\n    \n    'test_data_loader': {\n        'key': 'scenes/test.zarr',\n        'batch_size': 8,\n        'shuffle': False,\n        'num_workers': 4\n    }\n\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# set env variable for data\nos.environ[\"L5KIT_DATA_FOLDER\"] = DIR_INPUT\ndm = LocalDataManager(None)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model"},{"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, progress=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    # change the first layer input\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\n\n\ndef 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":{"trusted":true},"cell_type":"code","source":"criterion = nn.MSELoss(reduction=\"none\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\nmodel = build_model(cfg)\nmodel.to(device)\n\nmodel_state = torch.load(WEIGHT_FILE, map_location=device)\nmodel.load_state_dict(model_state)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Load Test Data"},{"metadata":{"trusted":true},"cell_type":"code","source":"# ===== INIT DATASET\ntest_cfg = cfg[\"test_data_loader\"]\n\n# Rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n\n# Test dataset/dataloader\ntest_zarr = ChunkedDataset(dm.require(test_cfg[\"key\"])).open()\ntest_mask = np.load(f\"{DIR_INPUT}/scenes/mask.npz\")[\"arr_0\"]\ntest_dataset = AgentDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\ntest_dataloader = DataLoader(test_dataset,\n                             shuffle=test_cfg[\"shuffle\"],\n                             batch_size=test_cfg[\"batch_size\"],\n                             num_workers=test_cfg[\"num_workers\"])\n\n\nprint(test_dataloader)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Predicting"},{"metadata":{"trusted":true},"cell_type":"code","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(test_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 = []\n    \n#     for agent_coords, world_from_agent, centroid in zip(agents_coords, world_from_agents, centroids):\n#         coords_offset.append(transform_points(agent_coords, world_from_agent) - centroid[: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())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model.eval()\n\nfuture_coords_offsets_pd = []\ntimestamps = []\nagent_ids = []\n\nwith torch.no_grad():\n    dataiter = tqdm(test_dataloader)\n    \n    for data in dataiter:\n        inputs = data[\"image\"].to(device)\n        target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(device)\n        targets = data[\"target_positions\"].to(device)\n\n        _, outputs = forward(data, model, device, criterion)\n        \n        future_coords_offsets_pd.append(outputs.cpu().numpy().copy())\n        timestamps.append(data[\"timestamp\"].numpy().copy())\n        agent_ids.append(data[\"track_id\"].numpy().copy())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"write_pred_csv('submission.csv',\n               timestamps=np.concatenate(timestamps),\n               track_ids=np.concatenate(agent_ids),\n               coords=np.concatenate(future_coords_offsets_pd))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","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}