{"cells":[{"metadata":{},"cell_type":"markdown","source":"Needless to say, one of the keys to this competition is the speed of data pre-processing.\nInspired by Peter's excellent code in the discussion below, I have created some rendering code in PyTorch3d.\n\nhttps://www.kaggle.com/c/lyft-motion-prediction-autonomous-vehicles/discussion/180359"},{"metadata":{},"cell_type":"markdown","source":"### This code currently suffers from the following problems\n* First of all, critically, the rendering is not fast enough.\n* We can't parallelize the pre-processing because processing gpu in the dataloader will cause CUDA initialization problems of Pytorch.\n* It doesn't include traffic light information. It is easy to solve this problem, but it has been postponed because the speed is not good enough."},{"metadata":{},"cell_type":"markdown","source":"If you have any methods of improvement, it would be very helpful if you could share them with me!"},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install pymap3d==2.1.0\n!pip install protobuf==3.12.2\n!pip install transforms3d\n!pip install zarr\n!pip install ptable\n!pip install --no-dependencies l5kit","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!conda install pytorch3d -c pytorch3d -y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%writefile tripy.py\n# MIT License\n#\n# Copyright (c) 2017 Sam Bolgert\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n#\n# https://github.com/linuxlewis/tripy\n#\nimport math\nimport sys\nfrom collections import namedtuple\n\nPoint = namedtuple('Point', ['idx', 'x', 'y', 'z']) # append point id\n\nEPSILON = math.sqrt(sys.float_info.epsilon)\n\n\ndef earclip(polygon):\n    \"\"\"\n    Simple earclipping algorithm for a given polygon p.\n    polygon is expected to be an array of 2-tuples of the cartesian points of the polygon\n    For a polygon with n points it will return n-2 triangles.\n    The triangles are returned as an array of 3-tuples where each item in the tuple is a 2-tuple of the cartesian point.\n    e.g\n    > polygon = [(0,1), (-1, 0), (0, -1), (1, 0)]\n    > triangles = tripy.earclip(polygon)\n    > triangles\n    [((1, 0), (0, 1), (-1, 0)), ((1, 0), (-1, 0), (0, -1))]\n    Implementation Reference:\n        - https://www.geometrictools.com/Documentation/TriangulationByEarClipping.pdf\n    \"\"\"\n    ear_vertex = []\n    triangles = []\n\n    polygon = [Point(idx, *point) for idx, point in enumerate(polygon)]\n\n    if _is_clockwise(polygon):\n        polygon.reverse()\n\n    point_count = len(polygon)\n    for i in range(point_count):\n        prev_index = i - 1\n        prev_point = polygon[prev_index]\n        point = polygon[i]\n        next_index = (i + 1) % point_count\n        next_point = polygon[next_index]\n\n        if _is_ear(prev_point, point, next_point, polygon):\n            ear_vertex.append(point)\n\n    while ear_vertex and point_count >= 3:\n        ear = ear_vertex.pop(0)\n        i = polygon.index(ear)\n        prev_index = i - 1\n        prev_point = polygon[prev_index]\n        next_index = (i + 1) % point_count\n        next_point = polygon[next_index]\n\n        polygon.remove(ear)\n        point_count -= 1\n        #triangles.append(((prev_point.x, prev_point.y, prev_point.z),\n        #                  (ear.x, ear.y, ear.z),\n        #                  (next_point.x, next_point.y, next_point.z)))\n        triangles.append((prev_point.idx, ear.idx, next_point.idx)) # append point id\n\n        if point_count > 3:\n            prev_prev_point = polygon[prev_index - 1]\n            next_next_index = (i + 1) % point_count\n            next_next_point = polygon[next_next_index]\n\n            groups = [\n                (prev_prev_point, prev_point, next_point, polygon),\n                (prev_point, next_point, next_next_point, polygon),\n            ]\n            for group in groups:\n                p = group[1]\n                if _is_ear(*group):\n                    if p not in ear_vertex:\n                        ear_vertex.append(p)\n                elif p in ear_vertex:\n                    ear_vertex.remove(p)\n    return triangles\n\n\ndef _is_clockwise(polygon):\n    s = 0\n    polygon_count = len(polygon)\n    for i in range(polygon_count):\n        point = polygon[i]\n        point2 = polygon[(i + 1) % polygon_count]\n        s += (point2.x - point.x) * (point2.y + point.y)\n    return s > 0\n\n\ndef _is_convex(prev, point, next_point):\n    return _triangle_sum(prev.x, prev.y, point.x, point.y, next_point.x, next_point.y) < 0\n\n\ndef _is_ear(p1, p2, p3, polygon):\n    ear = _contains_no_points(p1, p2, p3, polygon) and \\\n          _is_convex(p1, p2, p3) and \\\n          _triangle_area(p1.x, p1.y, p2.x, p2.y, p3.x, p3.y) > 0\n    return ear\n\n\ndef _contains_no_points(p1, p2, p3, polygon):\n    for pn in polygon:\n        if pn in (p1, p2, p3):\n            continue\n        elif _is_point_inside(pn, p1, p2, p3):\n            return False\n    return True\n\n\ndef _is_point_inside(p, a, b, c):\n    area = _triangle_area(a.x, a.y, b.x, b.y, c.x, c.y)\n    area1 = _triangle_area(p.x, p.y, b.x, b.y, c.x, c.y)\n    area2 = _triangle_area(p.x, p.y, a.x, a.y, c.x, c.y)\n    area3 = _triangle_area(p.x, p.y, a.x, a.y, b.x, b.y)\n    areadiff = abs(area - sum([area1, area2, area3])) < EPSILON\n    return areadiff\n\n\ndef _triangle_area(x1, y1, x2, y2, x3, y3):\n    return abs((x1 * (y2 - y3) + x2 * (y3 - y1) + x3 * (y1 - y2)) / 2.0)\n\n\ndef _triangle_sum(x1, y1, x2, y2, x3, y3):\n    return x1 * (y3 - y2) + x2 * (y1 - y3) + x3 * (y2 - y1)\n\n\ndef calculate_total_area(triangles):\n    result = []\n    for triangle in triangles:\n        sides = []\n        for i in range(3):\n            next_index = (i + 1) % 3\n            pt = triangle[i]\n            pt2 = triangle[next_index]\n            # Distance between two points\n            side = math.sqrt(math.pow(pt2[0] - pt[0], 2) + math.pow(pt2[1] - pt[1], 2))\n            sides.append(side)\n        # Heron's numerically stable forumla for area of a triangle:\n        # https://en.wikipedia.org/wiki/Heron%27s_formula\n        # However, for line-like triangles of zero area this formula can produce an infinitesimally negative value\n        # as an input to sqrt() due to the cumulative arithmetic errors inherent to floating point calculations:\n        # https://people.eecs.berkeley.edu/~wkahan/Triangle.pdf\n        # For this purpose, abs() is used as a reasonable guard against this condition.\n        c, b, a = sorted(sides)\n        area = .25 * math.sqrt(abs((a + (b + c)) * (c - (a - b)) * (c + (a - b)) * (a + (b - c))))\n        result.append((area, a, b, c))\n    triangle_area = sum(tri[0] for tri in result)\n    return triangle_area","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"%%writefile create_map_mesh.py\nimport pickle\nimport os\nfrom tqdm import tqdm\nimport numpy as np\n\nfrom l5kit.data.map_api import MapAPI\nfrom l5kit.data import ChunkedDataset, LocalDataManager\nfrom l5kit.rasterization.rasterizer_builder import (_load_metadata, get_hardcoded_world_to_ecef)\nimport transforms3d\n\nfrom tripy import earclip\n\ndef create_lane_surface(map_api):\n    verts = []\n    faces = []\n    count = 0\n    for element in tqdm(map_api):\n        element_id = MapAPI.id_as_str(element.id)\n        if map_api.is_lane(element):\n            lane = map_api.get_lane_coords(element_id)\n            left = lane[\"xyz_left\"]\n            right = lane[\"xyz_right\"]\n            points = np.vstack((left, np.flip(right, 0)))\n            triangles = np.array(earclip(points)) + count\n            verts.append(points)\n            faces.append(triangles)\n            count += points.shape[0]\n    return verts, faces\n\ndef create_crosswalks(map_api):\n    verts = []\n    faces = []\n    count = 0\n    for element in tqdm(map_api):\n        element_id = MapAPI.id_as_str(element.id)\n        if map_api.is_crosswalk(element):\n            crosswalk = map_api.get_crosswalk_coords(element_id)\n            points = crosswalk[\"xyz\"]\n            triangles = np.array(earclip(points)) + count\n            verts.append(points)\n            faces.append(triangles)\n            count += points.shape[0]\n    return verts, faces\n\ndef create_map_mesh():\n    filename = \"mesh.p\"\n\n    os.environ[\"L5KIT_DATA_FOLDER\"] = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n    dm = LocalDataManager(None)\n\n    # Rasterizer\n    semantic_map_path = dm.require(\"semantic_map/semantic_map.pb\")\n    dataset_meta = _load_metadata(\"meta.json\", dm)\n    world_to_ecef = np.array(dataset_meta[\"world_to_ecef\"], dtype=np.float64)\n\n    map_api = MapAPI(semantic_map_path, world_to_ecef)\n    lanes_verts, lanes_faces = create_lane_surface(map_api)\n    cw_verts, cw_faces = create_crosswalks(map_api)\n    with open(filename, \"wb\") as fp:\n        data = {\n            \"lane\": [lanes_verts, lanes_faces],\n            \"cw\": [cw_verts, cw_faces]\n        }\n        pickle.dump(data, fp)\n\nif __name__ == \"__main__\":\n    create_map_mesh()","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# preprocess: create map mesh data\n# !python create_map_mesh.py\n!ls /kaggle/input/lyft-map-mesh","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%writefile pytorch3d_rasterizer.py\nfrom typing import List, Optional, Tuple, Union\nimport pickle\nimport numpy as np\n\nfrom l5kit.data.filter import (filter_agents_by_labels, filter_agents_by_track_id)\nfrom l5kit.data.map_api import MapAPI\nfrom l5kit.geometry import world_to_image_pixels_matrix, transform_point, transform_points, yaw_as_rotation33\nfrom l5kit.rasterization import Rasterizer\nfrom l5kit.rasterization.box_rasterizer import get_ego_as_agent\n\nimport torch\nfrom pytorch3d.io import load_obj, save_obj\nfrom pytorch3d.structures import Meshes, join_meshes_as_batch, join_meshes_as_scene\nfrom pytorch3d.ops import sample_points_from_meshes\nfrom pytorch3d.transforms import Rotate, Translate\nfrom pytorch3d.renderer import (\n    FoVPerspectiveCameras, FoVOrthographicCameras,\n    look_at_view_transform, look_at_rotation, \n    RasterizationSettings, MeshRenderer, MeshRasterizer, BlendParams,\n    SoftSilhouetteShader, HardPhongShader, PointLights, TexturesVertex,\n    SoftPhongShader, HardFlatShader\n)\n\nimport transforms3d\n\nfrom tripy import earclip\n\n\ndef create_agetns(agents):\n    verts = []\n    faces = []\n    count = 0\n    corners_base_coords = np.asarray([[-1, -1], [-1, 1], [1, 1], [1, -1]])\n    for idx, agent in enumerate(agents):\n        corners = corners_base_coords * agent[\"extent\"][:2] / 2\n        r_m = yaw_as_rotation33(agent[\"yaw\"])\n        points = transform_points(corners, r_m) + agent[\"centroid\"][:2]\n        points = np.concatenate([points, np.ones([points.shape[0], 1]) * 0.03], axis=1)\n        triangles = np.array(earclip(points)) + count\n        verts.append(points)\n        faces.append(triangles)\n        count += points.shape[0]\n    return verts, faces\n\nclass PyTorch3dSemanticRasterizer(Rasterizer):\n\n    def __init__(self,\n                 raster_size: Tuple[int, int],\n                 pixel_size: Union[np.ndarray, list, float],\n                 ego_center: np.ndarray,\n                 filter_agents_threshold: float,\n                 history_num_frames: int,\n                 semantic_map_path: str,\n                 world_to_ecef: np.ndarray,\n                 filename: str):\n        super().__init__()\n\n        self.raster_size = raster_size\n        self.pixel_size = pixel_size\n        self.ego_center = ego_center\n\n        #if isinstance(pixel_size, np.ndarray) or isinstance(pixel_size, list):\n        #    self.pixel_size = pixel_size[0]\n\n        self.filter_agents_threshold = filter_agents_threshold\n\n        self.proto_API = MapAPI(semantic_map_path, world_to_ecef)\n\n        with open(filename, \"rb\") as fp:\n            data = pickle.load(fp)\n            lanes_verts, lanes_faces = data[\"lane\"]\n            cw_verts, cw_faces = data[\"cw\"]\n        \n        self.device = torch.device(\"cuda:0\")\n\n        # lanes mesh\n        lanes_verts_t = torch.Tensor(np.concatenate(lanes_verts, axis=0))\n        lanes_verts_t[:,2] = 1 # z axis\n        lanes_faces_t = torch.Tensor(np.concatenate(lanes_faces, axis=0))\n        lanes_verts_rgb = torch.zeros_like(lanes_verts_t)[None]  # (1, V, 3)\n        lanes_textures = TexturesVertex(verts_features=lanes_verts_rgb)\n        lanes_mesh = Meshes(verts=[lanes_verts_t], faces=[lanes_faces_t], textures=lanes_textures).to(self.device)\n\n        # crosswalks mesh\n        cw_verts_t = torch.Tensor(np.concatenate(cw_verts, axis=0))\n        cw_verts_t[:,2] = 2 # z axis\n        cw_faces_t = torch.Tensor(np.concatenate(cw_faces, axis=0))\n        cw_verts_rgb = torch.zeros_like(cw_verts_t)[None]  # (1, V, 3)\n        cw_verts_rgb[0,:] = torch.Tensor(np.array([255 / 255, 255 / 255, 0 / 255], dtype=np.float32))\n        cw_textures = TexturesVertex(verts_features=cw_verts_rgb)\n        cw_mesh = Meshes(verts=[cw_verts_t], faces=[cw_faces_t], textures=cw_textures).to(self.device)\n\n        # join mesh\n        self.map_mesh = join_meshes_as_scene([lanes_mesh, cw_mesh])\n\n        raster_settings = RasterizationSettings(\n            image_size=512,\n            blur_radius=0.0, \n            faces_per_pixel=1,\n        )\n\n        cameras = FoVPerspectiveCameras(device=self.device)\n\n        # lights = PointLights(device=device, location=((0.0, 0.0, -200.0),))\n\n        self.renderer = MeshRenderer(\n            rasterizer=MeshRasterizer(\n                cameras=cameras,\n                raster_settings=raster_settings\n            ),\n            # shader=SoftSilhouetteShader(device=device, cameras=cameras)\n            # shader=HardPhongShader(device=device, cameras=cameras, lights=lights)\n            shader=HardFlatShader(device=self.device, cameras=cameras)\n        )\n  \n    def rasterize(self, history_frames: np.ndarray, history_agents: List[np.ndarray],\n                  history_tl_faces: List[np.ndarray], agent: Optional[np.ndarray] = None,\n                  agents: Optional[np.ndarray] = None\n                  ) -> torch.Tensor:\n\n        with torch.no_grad():\n            frame = history_frames[0]\n            if agent is None:\n                translation = frame[\"ego_translation\"]\n                yaw = frame[\"ego_rotation\"]\n            else:\n                translation = agent[\"centroid\"]\n                yaw = agent[\"yaw\"]\n\n            world_to_image_space = world_to_image_pixels_matrix(\n                self.raster_size, self.pixel_size, translation, yaw, self.ego_center,\n            )\n            \n            center_pixel = np.asarray(self.raster_size) * (0.5, 0.5)\n            center_world = transform_point(center_pixel, np.linalg.inv(world_to_image_space))\n\n            for i, (frame, agents) in enumerate(zip(history_frames, history_agents)):\n                agents = filter_agents_by_labels(agents, self.filter_agents_threshold)\n                av_agent = get_ego_as_agent(frame).astype(agents.dtype)\n            \n                agents_verts, agents_faces = create_agetns(np.append(agents, av_agent))\n                agents_verts_t = torch.Tensor(np.concatenate(agents_verts, axis=0)).to(self.device, non_blocking=True)\n                agents_verts_t[:,2] = 3\n                agents_faces_t = torch.Tensor(np.concatenate(agents_faces, axis=0)).to(self.device, non_blocking=True)\n                agents_verts_rgb = torch.ones_like(agents_verts_t)[None].to(self.device, non_blocking=True)  # (1, V, 3)\n                agents_verts_rgb[0,:] = torch.Tensor(np.array([255 / 255, 0 / 255, 0 / 255], dtype=np.float32))\n                agents_textures = TexturesVertex(verts_features=agents_verts_rgb)\n                agents_mesh = Meshes(verts=[agents_verts_t], faces=[agents_faces_t], textures=agents_textures)\n\n                # no history frames\n                break\n\n            view_height = 225\n            camera_position = torch.Tensor(np.array([center_world[0], center_world[1], view_height])).to(self.device, non_blocking=True)\n            r = transforms3d.euler.euler2mat(0, np.pi, yaw)[None]\n            R = torch.Tensor(r).to(self.device, non_blocking=True)\n            T = -torch.bmm(R.transpose(1, 2), camera_position[None, :, None])[:, :, 0]\n            cameras = FoVPerspectiveCameras(device=self.device, R=R, T=T)\n\n            map_agents_mesh = join_meshes_as_scene([self.map_mesh, agents_mesh])\n\n            images = self.renderer(map_agents_mesh, cameras=cameras).squeeze(0)\n        return images\n\n    def to_rgb(self, in_im: np.ndarray, **kwargs: dict) -> np.ndarray:\n        return in_im","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%writefile ext_ego.py\nfrom typing import Optional, Tuple, cast\nimport warnings\n\nimport numpy as np\nimport torch\n\nfrom l5kit.data import DataManager, get_frames_slice_from_scenes\nfrom l5kit.dataset import AgentDataset\n\nclass ExtAgentDataset(AgentDataset):\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        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        frames = self.dataset.frames[get_frames_slice_from_scenes(self.dataset.scenes[scene_index])]\n\n        tl_faces = self.dataset.tl_faces\n        try:\n            if self.cfg[\"raster_params\"][\"disable_traffic_light_faces\"]:\n                tl_faces = np.empty(0, dtype=self.dataset.tl_faces.dtype)  # completely disable traffic light faces\n        except KeyError:\n            warnings.warn(\n                \"disable_traffic_light_faces not found in config, this will raise an error in the future\",\n                RuntimeWarning,\n                stacklevel=2,\n            )\n        data = self.sample_function(state_index, frames, self.dataset.agents, tl_faces, track_id)\n        # 0,1,C -> C,0,1\n        if isinstance(data[\"image\"], torch.Tensor):\n            image = data[\"image\"].permute(2, 0, 1)\n        else:\n            image = data[\"image\"].transpose(2, 0, 1)\n\n        target_positions = np.array(data[\"target_positions\"], 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_yaws = np.array(data[\"history_yaws\"], 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_yaws\": history_yaws,\n            \"history_availabilities\": data[\"history_availabilities\"],\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 os\nimport torch\n\nimport numpy as np\n%matplotlib inline\nimport matplotlib.pyplot as plt\n\nfrom torch import nn, optim\nfrom torch.utils.data import DataLoader\nfrom torchvision.models.resnet import resnet18\nfrom tqdm import tqdm\nfrom typing import Dict\n\nfrom l5kit.data import LocalDataManager, ChunkedDataset\nfrom l5kit.dataset import AgentDataset, EgoDataset\nfrom l5kit.rasterization import build_rasterizer","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"torch.cuda.is_available()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"from pytorch3d_rasterizer import PyTorch3dSemanticRasterizer\nfrom ext_ego import ExtAgentDataset","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DIR_INPUT = \"/kaggle/input/lyft-motion-prediction-autonomous-vehicles\"\n\nSINGLE_MODE_SUBMISSION = f\"{DIR_INPUT}/single_mode_sample_submission.csv\"\nMULTI_MODE_SUBMISSION = f\"{DIR_INPUT}/multi_mode_sample_submission.csv\"\n\nDEBUG = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\n        'model_architecture': 'resnet50',\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': [512, 512],\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    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 12,\n        'shuffle': True,\n        'num_workers': 0\n    },\n    \n    'sample_data_loader': {\n        'key': 'scenes/sample.zarr',\n        'batch_size': 12,\n        'shuffle': False,\n        'num_workers': 0\n    },\n    \n    'train_params': {\n        'max_num_steps': 100 if DEBUG else 10000,\n        'checkpoint_every_n_steps': 5000,\n        \n        # 'eval_every_n_steps': -1\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":"## pytorch3d rendering case"},{"metadata":{"trusted":true},"cell_type":"code","source":"from l5kit.rasterization.rasterizer_builder import (_load_metadata, get_hardcoded_world_to_ecef)\n\nraster_cfg = cfg[\"raster_params\"]\nsemantic_map_path = dm.require(raster_cfg[\"semantic_map_key\"])\ntry:\n    dataset_meta = _load_metadata(raster_cfg[\"dataset_meta_key\"], dm)\n    world_to_ecef = np.array(dataset_meta[\"world_to_ecef\"], dtype=np.float64)\nexcept (KeyError, FileNotFoundError):\n    world_to_ecef = get_hardcoded_world_to_ecef()\n\nrasterizer = PyTorch3dSemanticRasterizer(\n            raster_size=np.array(raster_cfg[\"raster_size\"]),\n            pixel_size=np.array(raster_cfg[\"pixel_size\"]),\n            ego_center=np.array(raster_cfg[\"ego_center\"]),\n            filter_agents_threshold=0.5,\n            history_num_frames=0,\n            semantic_map_path=semantic_map_path,\n            world_to_ecef=world_to_ecef,\n            filename=\"/kaggle/input/lyft-map-mesh/mesh.p\"\n        )\n\n# Train dataset/dataloader\ntrain_cfg = cfg[\"sample_data_loader\"]\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = ExtAgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset,\n                              shuffle=train_cfg[\"shuffle\"],\n                              batch_size=train_cfg[\"batch_size\"],\n                              num_workers=train_cfg[\"num_workers\"])\n\nprint(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfor data in train_dataloader:\n    print(data[\"image\"].shape)\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(8, 8))\nplt.imshow(data[\"image\"][0].permute(1,2,0).cpu().numpy()[::-1])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## original case"},{"metadata":{"trusted":true},"cell_type":"code","source":"cfg = {\n    'format_version': 4,\n    'model_params': {\n        'model_architecture': 'resnet50',\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': [512, 512],\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    'train_data_loader': {\n        'key': 'scenes/train.zarr',\n        'batch_size': 12,\n        'shuffle': True,\n        'num_workers': 2\n    },\n    \n    'sample_data_loader': {\n        'key': 'scenes/sample.zarr',\n        'batch_size': 12,\n        'shuffle': False,\n        'num_workers': 0\n    },\n    \n    'train_params': {\n        'max_num_steps': 100 if DEBUG else 10000,\n        'checkpoint_every_n_steps': 5000,\n        \n        # 'eval_every_n_steps': -1\n    }\n}","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# original case\n\n# Rasterizer\nrasterizer = build_rasterizer(cfg, dm)\n\n# Train dataset/dataloader\ntrain_cfg = cfg[\"sample_data_loader\"]\ntrain_zarr = ChunkedDataset(dm.require(train_cfg[\"key\"])).open()\ntrain_dataset = AgentDataset(cfg, train_zarr, rasterizer)\ntrain_dataloader = DataLoader(train_dataset,\n                              shuffle=train_cfg[\"shuffle\"],\n                              batch_size=train_cfg[\"batch_size\"],\n                              num_workers=train_cfg[\"num_workers\"])\n\nprint(train_dataset)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\nfor data in train_dataloader:\n    print(data[\"image\"].shape)\n    break","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"data[\"image\"].shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(8, 8))\nplt.imshow(data[\"image\"][0][22:].permute(1,2,0).cpu().numpy());\nplt.imshow(data[\"image\"][0][0].cpu().numpy(), cmap=\"Reds\", alpha=0.5);","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(8, 8))\nplt.imshow(data[\"image\"][0][22:].permute(1,2,0).cpu().numpy());","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}