{"cells":[{"metadata":{},"cell_type":"markdown","source":"<h1>Dataset class for use with Torch</h1>"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"!pip install -qqq lyft-dataset-sdk\n!pip install -qqq pyquaternion","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"import os\nimport cv2\nfrom pathlib import Path\nfrom functools import reduce\n\nimport numpy as np\nfrom pyquaternion import Quaternion\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nnp.random.seed(42)\ntorch.manual_seed(42)\n\nimport matplotlib.pyplot as plt\n\nfrom lyft_dataset_sdk.lyftdataset import LyftDataset\nfrom lyft_dataset_sdk.utils.data_classes import LidarPointCloud","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"input_dir = '/kaggle/input/3d-object-detection-for-autonomous-vehicles/'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!ln -sf {input_dir}/train_maps maps\n!ln -sf {input_dir}/train_images images\n!ln -sf {input_dir}/train_lidar lidar","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"I use the SDK for now since it makes it very easy to access all the properties. Might want to find a better way"},{"metadata":{"trusted":true},"cell_type":"code","source":"level5data = LyftDataset(\n    data_path='.',\n    json_path=os.path.join(input_dir + 'train_data'),\n    verbose=False\n)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h1>Utility functions</h1>"},{"metadata":{"trusted":true},"cell_type":"code","source":"def data2itemlist(data):\n    items = []\n    for scene in data.scene:\n        sample = data.get(\"sample\", scene[\"first_sample_token\"])\n        while sample:\n            data_dict = {k:data.get(\"sample_data\", v) for k, v in sample[\"data\"].items()}\n            for v in data_dict.values(): v[\"calibrated_sensor\"] = data.get(\"calibrated_sensor\", v[\"calibrated_sensor_token\"])\n            items.append({\n                **data_dict,\n                \"ego_pose\": data.get(\"ego_pose\", data_dict[\"CAM_FRONT\"][\"ego_pose_token\"]),\n                \"anns\": [data.get(\"sample_annotation\", x) for x in sample[\"anns\"]]\n            })\n            sample = data.get(\"sample\", sample[\"next\"]) if sample[\"next\"] else None\n    return items\n\ndef apply_rotation(coords, rotation):\n    q = Quaternion(rotation)\n    return np.dot(q.rotation_matrix, coords.T).T\n\ndef load_and_rotate(sensor):\n    data = LidarPointCloud.from_file(Path(sensor[\"filename\"])).points.T[:, :3]\n    calibration = sensor[\"calibrated_sensor\"]\n    data = apply_rotation(data, calibration[\"rotation\"]) + calibration[\"translation\"]\n    return data\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h1>Collate function for padding</h1>\n\nThis part needs to be thought out. For now it pads everything by adding zeros at the end (0 vectors at the end for LIDAR data, black pixels in the bottom/left part of images).\n\nNot only does every sample have a different number of LIDAR points, but the videos recorded by the cameras have a different shape.\nPossible solutions include cropping everything to the same size, or somehow unifying all the camera into some kind of 360 degrees view."},{"metadata":{"trusted":true},"cell_type":"code","source":"def collate_fn(samples):\n    data = {}\n    # list of dicts to dict of lists\n    samples = reduce(lambda x, y: {k:[v]+(x[k] if k in x else []) for k, v in y.items()}, samples, {})\n    for k, v in samples.items():\n        if v[0].shape:\n            max_length = [max([dp.shape[dim] for dp in v]) for dim in range(len(v[0].shape))]\n            padded = np.zeros(max_length)\n            for dp in v:\n                padded[tuple(slice(0, dp.shape[dim]) for dim in range(len(dp.shape)))] = dp\n            data[k] = torch.from_numpy(padded)\n        else:\n            data[k] = torch.from_numpy(np.array(v))\n    return data","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h1>Dataset class</h1>"},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftTorchDataset(Dataset):\n    def __init__(self, data, use_cache=False):\n        self.data = data\n        # To add or remove sensors\n        self.image_sensors = [\n            \"CAM_FRONT_RIGHT\",\n            \"CAM_FRONT\",\n            \"CAM_FRONT_LEFT\",\n            \"CAM_BACK_LEFT\",\n            \"CAM_BACK\",\n            \"CAM_BACK_RIGHT\",\n        ]\n        self.lidar_sensors = [\n            \"LIDAR_TOP\",\n        ]\n        self.items = [\n            item for item in data2itemlist(data)\n            if reduce(lambda x, y: x and (y in item), self.image_sensors + self.lidar_sensors, True)\n        ]\n        self.use_cache = use_cache\n        if use_cache:\n            self.cache_dict = {}\n    def __len__(self):\n        return len(self.items)\n    def __getitem__(self, index):\n        if self.use_cache and (index in self.cache):\n            return self.cache[index]    \n        image_data = {\n            sensor: cv2.imread(self.items[index][sensor][\"filename\"]).astype(np.float32)\n            for sensor in self.image_sensors\n        }\n        # TODO group by anchor?\n        lidar_data = {\n            sensor: load_and_rotate(self.items[index][sensor]).astype(np.float32)\n            for sensor in self.lidar_sensors\n        }\n        data = {\n            **image_data,\n            **lidar_data,\n            \"index\": np.array(index)\n        }\n        if self.use_cache:\n            self.cache[index] = data\n        return data","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Usage"},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = LyftTorchDataset(level5data)\ndataloader = DataLoader(dataset, batch_size=6, shuffle=True, collate_fn=collate_fn)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"example_batch = next(iter(dataloader))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"{k:v.shape for k, v in example_batch.items()}","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<h1>Stats</h1>\nLet's sample 100 random indeces and plot some stats on the data format"},{"metadata":{"trusted":true},"cell_type":"code","source":"shapes = np.array([(*dataset[idx][\"CAM_FRONT\"].shape, *dataset[idx][\"CAM_BACK\"].shape, *dataset[idx][\"LIDAR_TOP\"].shape) for idx in np.random.choice(len(dataset), size=100, replace=False)])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"plt.figure(figsize=(18, 4))\nplt.subplot(1, 5, 1)\nplt.title(\"CAM_FRONT height\")\nplt.hist(shapes[:,  0])\nplt.subplot(1, 5, 2)\nplt.title(\"CAM_FRONT width\")\nplt.hist(shapes[:, 1])\nplt.subplot(1, 5, 3)\nplt.title(\"CAM_BACK height\")\nplt.hist(shapes[:, 3])\nplt.subplot(1, 5, 4)\nplt.title(\"CAM_BACK height\")\nplt.hist(shapes[:, 4])\nplt.subplot(1, 5, 5)\nplt.title(\"LIDAR_TOP points\")\nplt.hist(shapes[:, 6])\nplt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"It seems about 80% of the data consists of 1224x1024 videos, and 20% of 1920x1080 videos. \nThere is significant overlap in the field of view recorded by CAM_FRONT and CAM_FRONT_RIGHT when they record in 1920x1080, but the overlap is quite small for 1224x1024.\nI assume they switched to 1224x1024 recording because their 45 degrees FOV is enough to capture 360 degrees. The 1920x1080 recordings can probably be cut to 1224x1024 without losing any information. I show an example below.\n\nWe see a similar distribution for the number of Lidar points, suggesting there's 2 different sensors that record each ~70000 points and ~110000 points.\nLidar points may be grouped into voxels with size that can be specified when creating the dataset.\nThese two techniques would produce tensors of fixed size, allowing us to use the default `collate_fn` function that doesn't apply any padding."},{"metadata":{"trusted":true},"cell_type":"code","source":"level5data.render_pointcloud_in_image(level5data.scene[0][\"first_sample_token\"], camera_channel=\"CAM_FRONT\")\nlevel5data.render_pointcloud_in_image(level5data.scene[0][\"first_sample_token\"], camera_channel=\"CAM_FRONT_RIGHT\")","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":1}