{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"from typing import List, Optional, Tuple\n\nimport numpy as np\n\nfrom l5kit.data import (\n    TL_FACE_DTYPE,\n    filter_agents_by_labels,\n    filter_tl_faces_by_frames,\n    get_agents_slice_from_frames,\n    get_tl_faces_slice_from_frames,\n)\nfrom l5kit.data.filter import filter_agents_by_frames, filter_agents_by_track_id\nfrom l5kit.geometry import rotation33_as_yaw, world_to_image_pixels_matrix\nfrom l5kit.kinematic import Perturbation\nfrom l5kit.rasterization import EGO_EXTENT_HEIGHT, EGO_EXTENT_LENGTH, EGO_EXTENT_WIDTH, Rasterizer\nfrom l5kit.sampling.slicing import get_future_slice, get_history_slice\n\n\ndef generate_kinetic_agent_sample(\n    state_index: int,\n    frames: np.ndarray,\n    agents: np.ndarray,\n    tl_faces: np.ndarray,\n    selected_track_id: Optional[int],\n    raster_size: Tuple[int, int],\n    pixel_size: np.ndarray,\n    ego_center: np.ndarray,\n    history_num_frames: int,\n    history_step_size: int,\n    future_num_frames: int,\n    future_step_size: int,\n    filter_agents_threshold: float,\n    rasterizer: Optional[Rasterizer] = None,\n    perturbation: Optional[Perturbation] = None,\n) -> dict:\n    \"\"\"Generates the inputs and targets to train a deep prediction model. A deep prediction model takes as input\n    the state of the world (here: an image we will call the \"raster\"), and outputs where that agent will be some\n    seconds into the future.\n\n    This function has a lot of arguments and is intended for internal use, you should try to use higher level classes\n    and partials that use this function.\n\n    Args:\n        state_index (int): The anchor frame index, i.e. the \"current\" timestep in the scene\n        frames (np.ndarray): The scene frames array, can be numpy array or a zarr array\n        agents (np.ndarray): The full agents array, can be numpy array or a zarr array\n        tl_faces (np.ndarray): The full traffic light faces array, can be numpy array or a zarr array\n        selected_track_id (Optional[int]): Either None for AV, or the ID of an agent that you want to\n        predict the future of. This agent is centered in the raster and the returned targets are derived from\n        their future states.\n        raster_size (Tuple[int, int]): Desired output raster dimensions\n        pixel_size (np.ndarray): Size of one pixel in the real world\n        ego_center (np.ndarray): Where in the raster to draw the ego, [0.5,0.5] would be the center\n        history_num_frames (int): Amount of history frames to draw into the rasters\n        history_step_size (int): Steps to take between frames, can be used to subsample history frames\n        future_num_frames (int): Amount of history frames to draw into the rasters\n        future_step_size (int): Steps to take between targets into the future\n        filter_agents_threshold (float): Value between 0 and 1 to use as cutoff value for agent filtering\n        based on their probability of being a relevant agent\n        rasterizer (Optional[Rasterizer]): Rasterizer of some sort that draws a map image\n        perturbation (Optional[Perturbation]): Object that perturbs the input and targets, used\nto train models that can recover from slight divergence from training set data\n\n    Raises:\n        ValueError: A ValueError is returned if the specified ``selected_track_id`` is not present in the scene\n        or was filtered by applying the ``filter_agent_threshold`` probability filtering.\n\n    Returns:\n        dict: a dict object with the raster array, the future offset coordinates (meters),\n        the future yaw angular offset, the future_availability as a binary mask\n    \"\"\"\n\n    agent_velocity = None\n\n    #  the history slice is ordered starting from the latest frame and goes backward in time., ex. slice(100, 91, -2)\n    history_slice = get_history_slice(state_index, history_num_frames, history_step_size, include_current_state=True)\n    future_slice = get_future_slice(state_index, future_num_frames, future_step_size)\n\n    history_frames = frames[history_slice].copy()  # copy() required if the object is a np.ndarray\n    future_frames = frames[future_slice].copy()\n\n    sorted_frames = np.concatenate((history_frames[::-1], future_frames))  # from past to future\n\n    # get agents (past and future)\n    agent_slice = get_agents_slice_from_frames(sorted_frames[0], sorted_frames[-1])\n    agents = agents[agent_slice].copy()  # this is the minimum slice of agents we need\n    history_frames[\"agent_index_interval\"] -= agent_slice.start  # sync interval with the agents array\n    future_frames[\"agent_index_interval\"] -= agent_slice.start  # sync interval with the agents array\n    history_agents = filter_agents_by_frames(history_frames, agents)\n    future_agents = filter_agents_by_frames(future_frames, agents)\n\n    try:\n        tl_slice = get_tl_faces_slice_from_frames(history_frames[-1], history_frames[0])  # -1 is the farthest\n        # sync interval with the traffic light faces array\n        history_frames[\"traffic_light_faces_index_interval\"] -= tl_slice.start\n        history_tl_faces = filter_tl_faces_by_frames(history_frames, tl_faces[tl_slice].copy())\n    except ValueError:\n        history_tl_faces = [np.empty(0, dtype=TL_FACE_DTYPE) for _ in history_frames]\n\n    if perturbation is not None:\n        history_frames, future_frames = perturbation.perturb(\n            history_frames=history_frames, future_frames=future_frames\n        )\n\n    # State you want to predict the future of.\n    cur_frame = history_frames[0]\n    cur_agents = history_agents[0]\n\n    if selected_track_id is None:\n        agent_centroid = cur_frame[\"ego_translation\"][:2]\n\n        # FIXME\n        # Assume velocity was 0, but it probably shouldnt be\n        agent_velocity = np.array([0, 0], dtype=np.float32)\n\n        agent_yaw = rotation33_as_yaw(cur_frame[\"ego_rotation\"])\n        agent_extent = np.asarray((EGO_EXTENT_LENGTH, EGO_EXTENT_WIDTH, EGO_EXTENT_HEIGHT))\n        selected_agent = None\n    else:\n        # this will raise IndexError if the agent is not in the frame or under agent-threshold\n        # this is a strict error, we cannot recover from this situation\n        try:\n            agent = filter_agents_by_track_id(\n                filter_agents_by_labels(cur_agents, filter_agents_threshold), selected_track_id\n            )[0]\n        except IndexError:\n            raise ValueError(f\" track_id {selected_track_id} not in frame or below threshold\")\n        agent_centroid = agent[\"centroid\"]\n        agent_velocity = agent[\"velocity\"]\n        agent_yaw = float(agent[\"yaw\"])\n        agent_extent = agent[\"extent\"]\n        selected_agent = agent\n\n    input_im = (\n        None\n        if not rasterizer\n        else rasterizer.rasterize(history_frames, history_agents, history_tl_faces, selected_agent)\n    )\n\n    world_to_image_space = world_to_image_pixels_matrix(\n        raster_size,\n        pixel_size,\n        ego_translation_m=agent_centroid,\n        ego_yaw_rad=agent_yaw,\n        ego_center_in_image_ratio=ego_center,\n    )\n\n    future_coords_offset, future_vel_offset, future_accel_offset, _, future_yaws_offset, future_availability = _create_targets_for_deep_prediction(\n        future_num_frames, future_frames, selected_track_id, future_agents, agent_centroid[:2], agent_velocity[:2],\n        agent_yaw,\n    )\n\n    # history_num_frames + 1 because it also includes the current frame\n    history_coords_offset, history_vel_offset, history_accel_offset, history_jerk_offset, history_yaws_offset, \\\n    history_availability = _create_targets_for_deep_prediction(history_num_frames + 1, history_frames,\n        selected_track_id, history_agents, agent_centroid[:2], agent_velocity[:2], agent_yaw,\n    )\n\n    # estimated_future_positions = history_jerk_offset\n    estimated_future_positions = _extrapolate_targets(future_num_frames, history_num_frames, history_coords_offset,\n                                                      history_vel_offset, history_accel_offset, history_jerk_offset)\n\n    error_arr = np.abs(estimated_future_positions - future_coords_offset[10:])\n    error = np.average(error_arr)\n\n    return {\n        \"image\": input_im,\n        \"target_positions\": future_coords_offset,\n        \"target_velocities\": future_vel_offset,\n        \"target_accelerations\": future_accel_offset,\n        \"target_yaws\": future_yaws_offset,\n        \"target_availabilities\": future_availability,\n        \"history_positions\": history_coords_offset,\n        \"history_velocities\": history_vel_offset,\n        \"history_accelerations\": history_accel_offset,\n        \"history_yaws\": history_yaws_offset,\n        \"history_availabilities\": history_availability,\n        \"estimated_future_positions\": estimated_future_positions,\n        \"world_to_image\": world_to_image_space,\n        \"centroid\": agent_centroid,\n        \"yaw\": agent_yaw,\n        \"extent\": agent_extent,\n    }\n\n\ndef _extrapolate_targets(future_num_frames: int, history_num_frames: int, position_history: np.ndarray,\n                         vel_history: np.ndarray, accel_history: np.ndarray, jerk_history: np.ndarray) -> np.ndarray:\n    \"\"\"\n    Extrapolates future positions using the first, second, and third derivatives of position\n\n    hehe not really extrapolation, but I'm just doing some estimation atm\n\n    :param future_num_frames:\n    :param position_history:\n    :param vel_history:\n    :param accel_history:\n    :param jerk_history:\n    :return: Extrapolated future positions - (future-history x 2)\n    \"\"\"\n\n    dt = 0.1\n\n    start_time = len(position_history) * dt\n\n    predict_frames = future_num_frames - history_num_frames\n\n    time_arr = np.arange(dt, predict_frames * dt + dt, dt)\n\n    avg_vel = np.array([np.average(vel_history.T[0]), np.average(vel_history.T[1])])\n    avg_acc = np.array([np.average(accel_history.T[0]), np.average(accel_history.T[1])])\n    avg_jrk = np.array([np.average(jerk_history.T[0]), np.average(jerk_history.T[1])])\n\n    pos = np.array([position_history[-1]]).T\n    vel = np.array([avg_vel]).T\n    accel = np.array([avg_acc]).T\n    jerk = np.array([avg_jrk]).T\n    # vel = np.array([vel_history[-1]]).T\n    # accel = np.array([accel_history[-1]]).T\n    # jerk = np.array([jerk_history[-1]]).T\n\n    # s =  \ts0 + v0t + ½a0t2 + ⅙jt3\n    res = pos * np.ones(time_arr.shape) + vel * time_arr + (accel * time_arr ** 2)/2 + (jerk * time_arr ** 3) / 6\n\n    return res.T\n\n\n\n\ndef _create_targets_for_deep_prediction(\n    num_frames: int,\n    frames: np.ndarray,\n    selected_track_id: Optional[int],\n    agents: List[np.ndarray],\n    agent_current_centroid: np.ndarray,\n    agent_current_velocity: np.ndarray,\n    agent_current_yaw: float,\n) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:\n    \"\"\"\n    Internal function that creates the targets and availability masks for deep prediction-type models.\n    The futures/history offset (in meters) are computed. When no info is available (e.g. agent not in frame)\n    a 0 is set in the availability array (1 otherwise).\n\n    Args:\n        num_frames (int): number of offset we want in the future/history\n        frames (np.ndarray): available frames. This may be less than num_frames\n        selected_track_id (Optional[int]): agent track_id or AV (None)\n        agents (List[np.ndarray]): list of agents arrays (same len of frames)\n        agent_current_centroid (np.ndarray): centroid of the agent at timestep 0\n        agent_current_yaw (float): angle of the agent at timestep 0\n\n    Returns:\n        Tuple[np.ndarray, np.ndarray, np.ndarray]: position offsets, angle offsets, availabilities\n\n    \"\"\"\n    #TODO: JERK\n\n    # Inclusively count the number of available and unavailable frames of this target\n    num_avail = 0\n    num_unavail = 0\n\n    last_avail_index = -1\n    last_unavail_index = -1\n\n    # How much the coordinates differ from the current state in meters.\n    coords_offset = np.zeros((num_frames, 2), dtype=np.float32)\n    vel_offset = np.zeros((num_frames, 2), dtype=np.float32)\n    accel_offset = np.zeros((num_frames, 2), dtype=np.float32)\n    jerk_offset = np.zeros((num_frames, 2), dtype=np.float32)\n    yaws_offset = np.zeros((num_frames, 1), dtype=np.float32)\n    availability = np.zeros((num_frames,), dtype=np.float32)\n\n    for i, (frame, agents) in enumerate(zip(frames, agents)):\n        if selected_track_id is None:\n            agent_centroid = frame[\"ego_translation\"][:2]\n            agent_yaw = rotation33_as_yaw(frame[\"ego_rotation\"])\n            agent_velocity = np.array([0, 0], dtype=np.float32)\n\n            num_avail = num_avail + 1\n        else:\n            # it's not guaranteed the target will be in every frame\n            try:\n                agent = filter_agents_by_track_id(agents, selected_track_id)[0]\n                num_avail = num_avail + 1\n            except IndexError:\n                availability[i] = 0.0  # keep track of invalid futures/history\n                last_unavail_index = i\n                num_unavail = num_unavail + 1\n                continue\n\n            agent_centroid = agent[\"centroid\"]\n            agent_velocity = agent[\"velocity\"]\n            agent_yaw = agent[\"yaw\"]\n\n        #TODO: change my kinetics to the format of\n        # offset for time[i] - time[0]\n        # So tidy up the velocity calculation process, and the acceleration stuff\n        # time 0 state, and time > 0 state. and make the calculations with respect to time passed and the time 0 val\n        # Also provide a means of getting change in time\n\n        # Note: B is an initial case\n        # State A: the boy scout case. All clean, no trouble\n        #   Two or more positions, previous frame is available\n        # State B: One or none positions\n        # State C: Two or more positions, previous frame is unavail\n\n        coords_offset[i] = agent_centroid - agent_current_centroid\n\n        # STATE B\n        if num_avail == 1:\n            # There is only the current position, no velocity is possible\n            vel_offset[i] = agent_current_velocity\n            accel_offset[i] = np.array([0, 0], dtype=np.float32)\n            jerk_offset[i] = np.array([0, 0], dtype=np.float32)\n\n        # there are two or more positions\n        # There has been no unavailabilities yet\n        elif last_unavail_index == -1:\n            # STATE A\n            # Only 1 frame difference in time\n            dx = coords_offset[i] - coords_offset[i - 1]\n            dt = 0.1\n\n            vel_offset[i] = dx / dt  # change in position / frame width (0.1 seconds) : [m/s]\n            # accel_offset[i] = (agent_velocity ** 2 - vel_offset[i - 1] ** 2) / (2 * dx)\n            accel_offset[i] = (vel_offset[i] - vel_offset[i - 1]) / dt\n            jerk_offset[i] = (accel_offset[i] - accel_offset[i - 1]) / dt\n\n        # Unavailability was previous frame\n        elif last_unavail_index == i - 1:\n            # Check number of positions\n            # STATE C\n            # There are positions before the unavailability, calculate the velocity\n            # Get previous existing velocity\n            prev_vel = vel_offset[last_avail_index]\n\n            # Get previous existing position\n            prev_pos = coords_offset[last_avail_index]\n\n            # Current Position\n            pos = coords_offset[i]\n\n            # Delta X\n            dx = pos - prev_pos\n\n            # Delta Time...0.1 being the frame width\n            dt = (i - last_avail_index) * 0.1\n\n            # Calculate velocity\n            vel = (dx / dt) * - prev_vel\n            vel_offset[i] = vel\n\n            # Now calculate acceleration\n            accel_offset[i] = (((dx - prev_vel * dt) * 2) ** (1/2)) / dt\n\n            # Now Jerk\n            jerk_offset[i] = (accel_offset[i] - accel_offset[last_avail_index]) / dt\n\n        else:\n            # STATE A\n            # Only 1 frame difference in time\n            dx = coords_offset[i] - coords_offset[i - 1]\n            dt = 0.1\n\n            vel_offset[i] = dx / dt  # change in position / frame width (0.1 seconds) : [m/s]\n            accel_offset[i] = (agent_velocity ** 2 - vel_offset[i - 1] ** 2) / (2 * dx)\n            jerk_offset[i] = (accel_offset[i] - accel_offset[i - 1]) / dt\n        # else:\n        #\n        #     # Attempt to get previous velocity\n        #     if i == 0:\n        #         # No previous velocity\n        #         prev_vel = np.array([0, 0], dtype=np.float32)\n        #     else:\n        #         # There must be a previous velocity\n        #         prev_vel = vel_offset[i - 1]\n        #\n        #     # Calculate the acceleration\n        #     accel_offset[i] = (vel_offset[i] ** 2 - prev_vel ** 2) / (2 * coords_offset[i])\n        # else:\n        #     accel_offset[i] = (agent_velocity ** 2 - agent_current_velocity ** 2) / (2 * coords_offset[i])\n\n        yaws_offset[i] = agent_yaw - agent_current_yaw\n        availability[i] = 1.0\n\n        # save this index as the last available index\n        last_avail_index = i\n    return coords_offset, vel_offset, accel_offset, jerk_offset, yaws_offset, availability","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from l5kit.dataset.ego import EgoDataset\nfrom l5kit.dataset.agent import AgentDataset\nfrom pathlib import Path\nimport bisect\n\nfrom l5kit.data import ChunkedDataset, get_agents_slice_from_frames, get_frames_slice_from_scenes\nfrom l5kit.kinematic import Perturbation\nfrom l5kit.rasterization import Rasterizer\nfrom l5kit.dataset.select_agents import TH_DISTANCE_AV, TH_EXTENT_RATIO, TH_YAW_DEGREE, select_agents\n\nimport numpy as np\nfrom zarr import convenience\nfrom functools import partial\nfrom typing import Optional, Tuple, cast\n\n\n# WARNING: changing these values impact the number of instances selected for both train and inference!\nMIN_FRAME_HISTORY = 10  # minimum number of frames an agents must have in the past to be picked\nMIN_FRAME_FUTURE = 1  # minimum number of frames an agents must have in the future to be picked\n\n\nclass KineticDataset(AgentDataset):\n\n    def __init__(\n        self,\n        cfg: dict,\n        zarr_dataset: ChunkedDataset,\n        rasterizer: Rasterizer,\n        perturbation: Optional[Perturbation] = None,\n        agents_mask: Optional[np.ndarray] = None,\n        min_frame_history: int = MIN_FRAME_HISTORY,\n        min_frame_future: int = MIN_FRAME_FUTURE,\n    ):\n        assert perturbation is None, \"KineticDataset does not support perturbation (yet)\"\n\n        super(KineticDataset, self).__init__(cfg, zarr_dataset, rasterizer, perturbation, agents_mask,\n                                             min_frame_history, min_frame_future)\n\n        # build a partial so we don't have to access cfg each time\n        self.sample_function = partial(\n            generate_kinetic_agent_sample,\n            raster_size=cast(Tuple[int, int], tuple(cfg[\"raster_params\"][\"raster_size\"])),\n            pixel_size=np.array(cfg[\"raster_params\"][\"pixel_size\"]),\n            ego_center=np.array(cfg[\"raster_params\"][\"ego_center\"]),\n            history_num_frames=cfg[\"model_params\"][\"history_num_frames\"],\n            history_step_size=cfg[\"model_params\"][\"history_step_size\"],\n            future_num_frames=cfg[\"model_params\"][\"future_num_frames\"],\n            future_step_size=cfg[\"model_params\"][\"future_step_size\"],\n            filter_agents_threshold=cfg[\"raster_params\"][\"filter_agents_threshold\"],\n            rasterizer=rasterizer,\n            perturbation=perturbation,\n        )\n\n    def get_frame(self, scene_index: int, state_index: int, track_id: Optional[int] = None) -> dict:\n        \"\"\"\n        A utility function to get the rasterisation and trajectory target for a given agent in a given frame\n\n        Args:\n            scene_index (int): the index of the scene in the zarr\n            state_index (int): a relative frame index in the scene\n            track_id (Optional[int]): the agent to rasterize or None for the AV\n        Returns:\n            dict: the rasterised image, the target trajectory (position and yaw) along with their availability,\n            the 2D matrix to center that agent, the agent track (-1 if ego) and the timestamp\n\n        \"\"\"\n        frames = self.dataset.frames[get_frames_slice_from_scenes(self.dataset.scenes[scene_index])]\n        data = self.sample_function(state_index, frames, self.dataset.agents, self.dataset.tl_faces, track_id)\n        # 0,1,C -> C,0,1\n        image = data[\"image\"].transpose(2, 0, 1)\n\n        target_positions = np.array(data[\"target_positions\"], dtype=np.float32)\n        # target_velocities = np.array(data[\"target_velocities\"], dtype=np.float32)\n        # target_accelerations = np.array(data[\"target_accelerations\"], dtype=np.float32)\n        target_yaws = np.array(data[\"target_yaws\"], dtype=np.float32)\n\n        history_positions = np.array(data[\"history_positions\"], dtype=np.float32)\n        # history_velocities = np.array(data[\"history_velocities\"], dtype=np.float32)\n        # history_accelerations = np.array(data[\"history_accelerations\"], dtype=np.float32)\n        # history_yaws = np.array(data[\"history_yaws\"], dtype=np.float32)\n\n        estimated_positions = np.array(data[\"estimated_future_positions\"], dtype=np.float32)\n\n        timestamp = frames[state_index][\"timestamp\"]\n        track_id = np.int64(-1 if track_id is None else track_id)  # always a number to avoid crashing torch\n\n        return {\n            \"image\": image,\n            \"target_positions\": target_positions,\n            \"target_yaws\": target_yaws,\n            \"target_availabilities\": data[\"target_availabilities\"],\n            \"history_positions\": history_positions,\n            \"history_availabilities\": data[\"history_availabilities\"],\n            \"estimated_future_positions\": estimated_positions,\n            \"world_to_image\": data[\"world_to_image\"],\n            \"track_id\": track_id,\n            \"timestamp\": timestamp,\n            \"centroid\": data[\"centroid\"],\n            \"yaw\": data[\"yaw\"],\n            \"extent\": data[\"extent\"],\n        }","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nfrom torchvision.models import mnasnet1_0\n\n\ndef trim_network_at_index(network: nn.Module, index: int = -1) -> nn.Module:\n    \"\"\"\n    Returns a new network with all layers up to index from the back.\n    :param network: Module to trim.\n    :param index: Where to trim the network. Counted from the last layer.\n    \"\"\"\n    assert index < 0, f\"Param index must be negative. Received {index}.\"\n    return nn.Sequential(*list(network.children())[:index])\n\n\nclass MnasBackbone(nn.Module):\n    \"\"\"\n    Outputs tensor after last convolution before the fully connected layer.\n\n    Allowed versions: mobilenet_v2.\n    \"\"\"\n\n    def __init__(self, num_in_channels: int = 3):\n        \"\"\"\n        Inits MobileNetBackbone.\n        :param version: mobilenet version to use.\n        \"\"\"\n        super().__init__()\n\n        model = mnasnet1_0(pretrained=False)\n        layer = model.layers[0]\n        model.layers[0] = nn.Conv2d(\n            in_channels=num_in_channels,\n            out_channels=layer.out_channels,\n            kernel_size=layer.kernel_size,\n            stride=layer.stride,\n            padding=layer.padding,\n            bias=False,\n        )\n\n        self.backbone = trim_network_at_index(model, -1)\n\n    def forward(self, input_tensor: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Outputs features after last convolution.\n        :param input_tensor:  Shape [batch_size, n_channels, length, width].\n        :return: Tensor of shape [batch_size, n_convolution_filters]. For mobilenet_v2,\n            the shape is [batch_size, 1280].\n        \"\"\"\n        backbone_features = self.backbone(input_tensor)\n        return backbone_features.mean([2, 3])\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import math\nimport random\nfrom typing import List, Tuple\n\nimport torch\nfrom torch import nn\nfrom torch.nn import functional as f\n\n# Number of entries in Agent State Vector\nASV_DIM = 3\n\ndef calculate_backbone_feature_dim(backbone, input_shape: Tuple[int, int, int]) -> int:\n    \"\"\" Helper to calculate the shape of the fully-connected regression layer. \"\"\"\n    tensor = torch.ones(1, *input_shape)\n    output_feat = backbone.forward(tensor)\n    return output_feat.shape[-1]\n\nclass LyftNet(nn.Module):\n    \"\"\"\n    Implementation of Multiple-Trajectory Prediction (MTP) model\n    based on https://arxiv.org/pdf/1809.10732.pdf\n    \"\"\"\n\n    def __init__(self, backbone: nn.Module, num_modes: int, num_targets: int, num_kinetic_dim: int = 3,\n                 n_hidden_layers: int = 4096, input_shape: Tuple[int, int, int] = (3, 500, 500)):\n        \"\"\"\n        Inits the MTP network.\n        :param backbone: CNN Backbone to use.\n        :param num_modes: Number of predicted paths to estimate for each agent.\n        :param seconds: Number of seconds into the future to predict.\n            Default for the challenge is 6.\n        :param frequency_in_hz: Frequency between timesteps in the prediction (in Hz).\n            Highest frequency is nuScenes is 2 Hz.\n        :param n_hidden_layers: Size of fully connected layer after the CNN\n            backbone processes the image.\n        :param input_shape: Shape of the input expected by the network.\n            This is needed because the size of the fully connected layer after\n            the backbone depends on the backbone and its version.\n        Note:\n            Although seconds and frequency_in_hz are typed as floats, their\n            product should be an int.\n        \"\"\"\n\n        super().__init__()\n\n        self.ASV_DIM = num_kinetic_dim\n\n        self.backbone = backbone\n        self.num_modes = num_modes\n        backbone_feature_dim = calculate_backbone_feature_dim(backbone, input_shape)\n        self.fc1 = nn.Linear(backbone_feature_dim + self.ASV_DIM, n_hidden_layers, bias=True)\n        # predictions_per_mode = int(seconds * frequency_in_hz) * 2\n        predictions_per_mode = num_targets\n\n        self.num_pred = predictions_per_mode * num_modes\n        self.future_len = int(num_targets / 2)\n\n        self.norm = nn.BatchNorm1d(self.ASV_DIM)\n\n        self.fc2 = nn.Linear(n_hidden_layers, int(self.num_pred + num_modes))\n\n    def forward(self, image_tensor: torch.Tensor,\n                agent_state_vector: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Forward pass of the model.\n        :param image_tensor: Tensor of images shape [batch_size, n_channels, length, width].\n        :param agent_state_vector: Tensor of floats representing the agent state.\n            [batch_size, ASV_DIM * num_targets].\n        :returns: Tensor of dimension [batch_size, number_of_modes * number_of_predictions_per_mode + number_of_modes]\n            storing the predicted trajectory and mode probabilities. Mode probabilities are normalized to sum\n            to 1 during inference.\n        \"\"\"\n\n        backbone_features = self.backbone(image_tensor)\n\n        features = torch.cat([backbone_features, self.norm(agent_state_vector)], dim=1)\n\n        predictions = self.fc2(self.fc1(features))\n\n        bs, _ = predictions.shape\n\n        pred, confidences = torch.split(predictions, self.num_pred, dim=1)\n\n        pred = pred.view(bs, self.num_modes, self.future_len, 2)\n        assert confidences.shape == (bs, self.num_modes)\n\n        confidences = torch.softmax(confidences, dim=1)\n\n        # print(self.fc1.weight)\n        # print(\"---------------\")\n\n        # nn.utils.clip_grad_norm_()\n\n        if math.isnan(confidences[0][0].item()):\n            print(\"Oh no, not the naan\")\n            print(agent_state_vector)\n\n        return pred, confidences","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# --- Function utils ---\n# Original code from https://github.com/lyft/l5kit/blob/20ab033c01610d711c3d36e1963ecec86e8b85b6/l5kit/l5kit/evaluation/metrics.py\nimport numpy as np\n\nimport torch\nfrom torch import Tensor\n\ndef pytorch_neg_multi_log_likelihood_batch(\n    gt: Tensor, pred: Tensor, confidences: Tensor, avails: Tensor\n) -> Tensor:\n    \"\"\"\n    Compute a negative log-likelihood for the multi-modal scenario.\n    log-sum-exp trick is used here to avoid underflow and overflow, For more information about it see:\n    https://en.wikipedia.org/wiki/LogSumExp#log-sum-exp_trick_for_log-domain_calculations\n    https://timvieira.github.io/blog/post/2014/02/11/exp-normalize-trick/\n    https://leimao.github.io/blog/LogSumExp/\n    Args:\n        gt (Tensor): array of shape (bs)x(time)x(2D coords)\n        pred (Tensor): array of shape (bs)x(modes)x(time)x(2D coords)\n        confidences (Tensor): array of shape (bs)x(modes) with a confidence for each mode in each sample\n        avails (Tensor): array of shape (bs)x(time) with the availability for each gt timestep\n    Returns:\n        Tensor: negative log-likelihood for this example, a single float number\n    \"\"\"\n    assert len(pred.shape) == 4, f\"expected 3D (MxTxC) array for pred, got {pred.shape}\"\n    batch_size, num_modes, future_len, num_coords = pred.shape\n\n    assert gt.shape == (batch_size, future_len, num_coords), f\"expected 2D (Time x Coords) array for gt, got {gt.shape}\"\n    assert confidences.shape == (batch_size, num_modes), f\"expected 1D (Modes) array for gt, got {confidences.shape}\"\n    if not torch.allclose(torch.sum(confidences, dim=1), confidences.new_ones((batch_size,))):\n        print(confidences)\n    assert torch.allclose(torch.sum(confidences, dim=1), confidences.new_ones((batch_size,))), \"confidences should sum to 1\"\n    assert avails.shape == (batch_size, future_len), f\"expected 1D (Time) array for gt, got {avails.shape}\"\n    # assert all data are valid\n    assert torch.isfinite(pred).all(), \"invalid value found in pred\"\n    assert torch.isfinite(gt).all(), \"invalid value found in gt\"\n    assert torch.isfinite(confidences).all(), \"invalid value found in confidences\"\n    assert torch.isfinite(avails).all(), \"invalid value found in avails\"\n\n    # convert to (batch_size, num_modes, future_len, num_coords)\n    gt = torch.unsqueeze(gt, 1)  # add modes\n    avails = avails[:, None, :, None]  # add modes and cords\n\n    # error (batch_size, num_modes, future_len)\n    error = torch.sum(((gt - pred) * avails) ** 2, dim=-1)  # reduce coords and use availability\n\n    with np.errstate(divide=\"ignore\"):  # when confidence is 0 log goes to -inf, but we're fine with it\n        # error (batch_size, num_modes)\n        error = torch.log(confidences) - 0.5 * torch.sum(error, dim=-1)  # reduce time\n\n    # use max aggregator on modes for numerical stability\n    # error (batch_size, num_modes)\n    max_value, _ = error.max(dim=1, keepdim=True)  # error are negative at this point, so max() gives the minimum one\n    error = -torch.log(torch.sum(torch.exp(error - max_value), dim=-1, keepdim=True)) - max_value  # reduce modes\n    # print(\"error\", error)\n    return torch.mean(error)\n\n\ndef pytorch_neg_multi_log_likelihood_single(\n    gt: Tensor, pred: Tensor, avails: Tensor\n) -> Tensor:\n    \"\"\"\n\n    Args:\n        gt (Tensor): array of shape (bs)x(time)x(2D coords)\n        pred (Tensor): array of shape (bs)x(time)x(2D coords)\n        avails (Tensor): array of shape (bs)x(time) with the availability for each gt timestep\n    Returns:\n        Tensor: negative log-likelihood for this example, a single float number\n    \"\"\"\n    # pred (bs)x(time)x(2D coords) --> (bs)x(mode=1)x(time)x(2D coords)\n    # create confidence (bs)x(mode=1)\n    batch_size, future_len, num_coords = pred.shape\n    confidences = pred.new_ones((batch_size, 1))\n    return pytorch_neg_multi_log_likelihood_batch(gt, pred.unsqueeze(1), confidences, avails)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from l5kit.evaluation import write_pred_csv\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.rasterization import build_rasterizer\nfrom l5kit.dataset import AgentDataset\n\nfrom torch.utils.data import DataLoader\n\nimport os\nimport torch\n\nfrom tqdm import tqdm\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\n\nimport gc\n\nimport pandas as pd\n\n\nclass LyftManager:\n\n    def __init__(self, config, data_path, device, num_modes=3, verbose=False):\n\n        self.cfg = config\n        self.data_path = data_path\n        self.device = device\n\n        self.verbose = verbose\n\n        # change input channels number to match the rasterizer's output\n        num_history_channels = (self.cfg[\"model_params\"][\"history_num_frames\"] + 1) * 2\n        num_in_channels = 3 + num_history_channels\n\n        # Create backbone\n        self.backbone = MnasBackbone(num_in_channels=num_in_channels)\n\n        # Raster Size\n        raster_size = self.cfg[\"raster_params\"][\"raster_size\"]\n\n        # change output size to (X, Y) * number of future states\n        num_targets = 2 * self.cfg[\"model_params\"][\"future_num_frames\"]\n        self.future_len = self.cfg[\"model_params\"][\"future_num_frames\"]\n        print(\"Number of Targets = \", self.future_len)\n        self.num_targets = num_targets\n        self.num_preds = num_targets * num_modes\n        self.num_modes = num_modes\n\n        self.model = LyftNet(self.backbone, num_modes=num_modes, num_kinetic_dim=2 * (self.future_len + 1),\n                             num_targets=num_targets, input_shape=(num_in_channels, raster_size[0], raster_size[1]))\n\n        self.model.to(device=self.device)\n\n    def train(self, iterations, lr=1e-3, file_name=\"lyft-net.pth\"):\n\n        # set env variable for data\n        os.environ[\"L5KIT_DATA_FOLDER\"] = self.data_path\n        dm = LocalDataManager(None)\n        # get config\n        cfg = self.cfg\n        print(cfg)\n\n        # ===== INIT DATASET\n        train_cfg = cfg[\"train_data_loader\"]\n\n        # Rasterizer\n        rasterizer = build_rasterizer(cfg, dm)\n\n        # Train dataset/dataloader\n        train_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\n        if self.verbose:\n            print(\"Dataset Chunked\")\n        train_dataset = KineticDataset(cfg, train_zarr, rasterizer)\n        if self.verbose:\n            print(\"Agent Dataset Retrieved\")\n\n        # agents = pd.DataFrame.from_records(train_zarr.agents,\n        #                                    columns=['centroid', 'extent', 'yaw', 'velocity', 'track_id',\n        #                                             'label_probabilities'])\n\n        #\n        # print(agents.head())\n\n        train_dataloader = DataLoader(train_dataset,\n                                      shuffle=train_cfg[\"shuffle\"],\n                                      batch_size=train_cfg[\"batch_size\"],\n                                      num_workers=train_cfg[\"num_workers\"])\n\n        print(train_dataset)\n\n        # ==== INIT MODEL parameters\n        optimizer = optim.Adam(self.model.parameters(), lr=lr)\n        criterion = nn.PoissonNLLLoss()\n\n        # ==== TRAIN LOOP\n\n        tr_it = iter(train_dataloader)\n        progress_bar = tqdm(range(iterations))\n        losses_train = []\n        rolling_avg = []\n        # torch.save(model.state_dict(), \"/home/michael/Workspace/Lyft/model/resnet_base.pth\")\n        for i 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            self.model.train()\n            torch.set_grad_enabled(True)\n\n            ####################\n            # this is where it gets real funky\n\n            # ids = data[\"track_id\"]\n            position_tensor = data[\"target_positions\"].to(self.device)\n            # velocity_tensor = data[\"target_velocities\"].to(self.device)\n            # acceleration_tensor = data[\"target_accelerations\"].to(self.device)\n            # yaw_tensor = data[\"target_yaws\"].to(self.device)\n\n            history_position_tensor = data[\"history_positions\"].to(self.device)\n            estimated_future_positions = data[\"estimated_future_positions\"].to(self.device)\n            # history_velocity_tensor = data[\"history_velocities\"].to(self.device)\n            # history_acceleration_tensor = data[\"history_accelerations\"].to(self.device)\n            # history_yaw_tensor = data[\"history_yaws\"].to(self.device)\n            # history_availability = data[\"history_availabilities\"].to(self.device)\n\n            imageTensor = data[\"image\"].to(self.device)\n            if self.verbose:\n                print(\"Image Tensor: \", imageTensor.shape)\n\n            # state_vector = torch.cat([history_position_tensor, history_velocity_tensor, history_acceleration_tensor,\n            #                           history_yaw_tensor], 2).to(self.device)\n\n            state_vector = torch.cat([estimated_future_positions, history_position_tensor], 1).to(self.device)\n\n            state_vector = torch.flatten(state_vector, 1).to(self.device)\n            if self.verbose:\n                print(\"State Vector: \", state_vector.shape)\n\n            pred, conf = self.model.forward(imageTensor, state_vector)\n\n            # loss2 = criterion(pred, position_tensor)\n\n            # print(loss2)\n\n            if self.verbose:\n                print(\"Prediction: \", pred.shape)\n                print(\"Probability: \", conf.shape)\n\n            target_availabilities = data[\"target_availabilities\"].unsqueeze(-1).to(self.device)\n\n            flattenedTargets = torch.flatten(target_availabilities, 1, 2)\n\n            loss = pytorch_neg_multi_log_likelihood_batch(position_tensor, pred, conf, flattenedTargets)\n\n            # loss = self.lossModel(pred, data[\"target_positions\"])\n\n\n            # Backward pass\n            optimizer.zero_grad()\n            loss.backward()\n            optimizer.step()\n\n            losses_train.append(loss.item())\n            rolling_avg.append(np.mean(losses_train))\n            progress_bar.set_description(f\"loss: {loss.item()} loss(avg): {np.mean(losses_train)}\")\n\n            # if i == 10000:\n            #     torch.save(model.state_dict(), \"/home/michael/Workspace/Lyft/model/resnet\" + str(i) + \".pth\")\n\n        print(\"Done Training\")\n        torch.save(self.model.state_dict(), f\"/home/michael/Workspace/Lyft/model/{file_name}\")\n\n    def evaluate(self, data_path, file_name=\"submission.csv\"):\n\n        # set env variable for data\n        os.environ[\"L5KIT_DATA_FOLDER\"] = data_path\n        dm = LocalDataManager(None)\n\n        cfg = self.cfg\n\n        # ===== INIT DATASET\n        test_cfg = cfg[\"test_data_loader\"]\n\n        # Rasterizer\n        rasterizer = build_rasterizer(cfg, dm)\n\n        # Test dataset/dataloader\n        test_zarr = ChunkedDataset(dm.require(test_cfg[\"key\"])).open()\n        test_mask = np.load(f\"{data_path}/scenes/mask.npz\")[\"arr_0\"]\n        # test_dataset = AgentDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\n        test_dataset = KineticDataset(cfg, test_zarr, rasterizer, agents_mask=test_mask)\n        test_dataloader = DataLoader(test_dataset,\n                                     shuffle=test_cfg[\"shuffle\"],\n                                     batch_size=test_cfg[\"batch_size\"],\n                                     num_workers=test_cfg[\"num_workers\"])\n        test_dataloader = test_dataloader\n        print(test_dataloader)\n\n        # ==== EVAL LOOP\n        self.model.eval()\n        torch.set_grad_enabled(False)\n        criterion = nn.MSELoss(reduction=\"none\")\n\n        # store information for evaluation\n        future_coords_offsets_pd = []\n        timestamps = []\n        pred_coords =  []\n        confidences_list = []\n\n        agent_ids = []\n        progress_bar = tqdm(test_dataloader)\n        for data in progress_bar:\n\n            # ids = data[\"track_id\"]\n            # position_tensor = data[\"target_positions\"].to(self.device)\n            # velocity_tensor = data[\"target_velocities\"].to(self.device)\n            # acceleration_tensor = data[\"target_accelerations\"].to(self.device)\n            # yaw_tensor = data[\"target_yaws\"].to(self.device)\n\n            history_position_tensor = data[\"history_positions\"].to(self.device)\n            estimated_future_positions = data[\"estimated_future_positions\"].to(self.device)\n            # history_velocity_tensor = data[\"history_velocities\"].to(self.device)\n            # history_acceleration_tensor = data[\"history_accelerations\"].to(self.device)\n            # history_yaw_tensor = data[\"history_yaws\"].to(self.device)\n            # history_availability = data[\"history_availabilities\"].to(self.device)\n\n            imageTensor = data[\"image\"].to(self.device)\n            if self.verbose:\n                print(\"Image Tensor: \", imageTensor.shape)\n\n            # state_vector = torch.cat([history_position_tensor, history_velocity_tensor, history_acceleration_tensor,\n            #                           history_yaw_tensor], 2).to(self.device)\n\n            state_vector = torch.cat([estimated_future_positions, history_position_tensor], 1).to(self.device)\n\n            state_vector = torch.flatten(state_vector, 1).to(self.device)\n            # print(state_vector)\n            if self.verbose:\n                print(\"State Vector: \", state_vector.shape)\n\n            pred, confidences = self.model.forward(imageTensor, state_vector)\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())\n            #\n            # pred, confidences = predictor(image)\n\n            pred_coords.append(pred.cpu().numpy().copy())\n            confidences_list.append(confidences.cpu().numpy().copy())\n\n        # ==== Save Results\n        pred_path = f\"{os.getcwd()}/{file_name}\"\n        write_pred_csv(pred_path,\n                       timestamps=np.concatenate(timestamps),\n                       track_ids=np.concatenate(agent_ids),\n                       coords=np.concatenate(pred_coords),\n                       confs=np.concatenate(confidences_list))\n\n        torch.cuda.empty_cache()\n\n    def load(self, checkpoint_path:str):\n\n        state_dict = torch.load(checkpoint_path, map_location=self.device)\n        self.model.load_state_dict(state_dict)\n\n    # def clean(self):  # DOES WORK\n    #     self._optimizer_to(torch.device('cuda:0'))\n    #     del self.optimizer\n    #     gc.collect()\n    #     torch.cuda.empty_cache()\n    #\n    # def _optimizer_to(self, device):\n    #     for param in self.optimizer.state.values():\n    #         # Not sure there are any global tensors in the state dict\n    #         if isinstance(param, torch.Tensor):\n    #             param.data = param.data.to(device)\n    #             if param._grad is not None:\n    #                 param._grad.data = param._grad.data.to(device)\n    #         elif isinstance(param, dict):\n    #             for subparam in param.values():\n    #                 if isinstance(subparam, torch.Tensor):\n    #                     subparam.data = subparam.data.to(device)\n    #                     if subparam._grad is not None:\n    #                         subparam._grad.data = subparam._grad.data.to(device)\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nfrom l5kit.configs import load_config_data\n\ndata_path = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n\ndevice = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\nconfig = load_config_data(\"/kaggle/input/lyftnet/config_lyftnet.yaml\")\n\nconfig[\"test_data_loader\"][\"batch_size\"] = 16\nconfig[\"test_data_loader\"][\"num_workers\"] = 8\n\ncheckpoint_path = \"/kaggle/input/lyftnet/lyft-net.pth\"\n\n\nmanager = LyftManager(config=config, data_path=data_path, device=device, num_modes=3, verbose=False)\n\nmanager.load(checkpoint_path)\n\n#toto train for 324000\n# manager.train(iterations=44000, lr=1e-4, file_name=\"lyft-net.pth\")\n# manager.train(iterations=324000, lr=1e-4, file_name=\"lyft-net.pth\")\n\nmanager.evaluate(data_path=data_path)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# import os\n# os.system(\"!cp /kaggle/input/lyftnet/submission.csv /kaggle/working/submission.csv\")","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}