{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":15768,"databundleVersionId":700263,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install shapely \n!pip install lyft-dataset-sdk","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:39:45.318382Z","iopub.execute_input":"2024-09-10T13:39:45.318796Z","iopub.status.idle":"2024-09-10T13:40:20.704198Z","shell.execute_reply.started":"2024-09-10T13:39:45.318756Z","shell.execute_reply":"2024-09-10T13:40:20.703113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Our code will generate data, visualization and model checkpoints, they will be persisted to disk in this folder\nARTIFACTS_FOLDER = \"./artifacts\"","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:41:50.571178Z","iopub.execute_input":"2024-09-10T13:41:50.572269Z","iopub.status.idle":"2024-09-10T13:41:50.576438Z","shell.execute_reply.started":"2024-09-10T13:41:50.572225Z","shell.execute_reply":"2024-09-10T13:41:50.575552Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datetime import datetime\nfrom functools import partial\nimport glob\nfrom multiprocessing import Pool\n\n# Disable multiprocesing for numpy/opencv. We already multiprocess ourselves, this would mean every subprocess produces\n# even more threads which would lead to a lot of context switching, slowing things down a lot.\nimport os\nos.environ[\"OMP_NUM_THREADS\"] = \"1\"\n\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nimport pandas as pd\nimport cv2\nfrom PIL import Image\nimport numpy as np\nfrom tqdm import tqdm, tqdm_notebook\nimport scipy\nimport scipy.ndimage\nimport scipy.special\nfrom scipy.spatial.transform import Rotation as R\n\nfrom lyft_dataset_sdk.lyftdataset import LyftDataset\nfrom lyft_dataset_sdk.utils.data_classes import LidarPointCloud, Box, Quaternion\nfrom lyft_dataset_sdk.utils.geometry_utils import view_points, transform_matrix","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:41:52.665492Z","iopub.execute_input":"2024-09-10T13:41:52.665908Z","iopub.status.idle":"2024-09-10T13:41:52.676906Z","shell.execute_reply.started":"2024-09-10T13:41:52.665867Z","shell.execute_reply":"2024-09-10T13:41:52.67598Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!ln -s /kaggle/input/3d-object-detection-for-autonomous-vehicles/train_images images\n!ln -s /kaggle/input/3d-object-detection-for-autonomous-vehicles/train_maps maps\n!ln -s /kaggle/input/3d-object-detection-for-autonomous-vehicles/train_lidar lidar","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:41:59.270343Z","iopub.execute_input":"2024-09-10T13:41:59.270735Z","iopub.status.idle":"2024-09-10T13:42:02.444847Z","shell.execute_reply.started":"2024-09-10T13:41:59.270698Z","shell.execute_reply":"2024-09-10T13:42:02.443412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"level5data = LyftDataset(data_path='.', json_path='/kaggle/input/3d-object-detection-for-autonomous-vehicles/train_data', verbose=True)\nos.makedirs(ARTIFACTS_FOLDER, exist_ok=True)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:42:06.31834Z","iopub.execute_input":"2024-09-10T13:42:06.318775Z","iopub.status.idle":"2024-09-10T13:42:30.173272Z","shell.execute_reply.started":"2024-09-10T13:42:06.318737Z","shell.execute_reply":"2024-09-10T13:42:30.172244Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"classes = [\"car\", \"motorcycle\", \"bus\", \"bicycle\", \"truck\", \"pedestrian\", \"other_vehicle\", \"animal\", \"emergency_vehicle\"]","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:43:56.932443Z","iopub.execute_input":"2024-09-10T13:43:56.933213Z","iopub.status.idle":"2024-09-10T13:43:56.938066Z","shell.execute_reply.started":"2024-09-10T13:43:56.933168Z","shell.execute_reply":"2024-09-10T13:43:56.936955Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"records = [(level5data.get('sample', record['first_sample_token'])['timestamp'], record) for record in level5data.scene]\n\nentries = []\n\nfor start_time, record in sorted(records):\n    start_time = level5data.get('sample', record['first_sample_token'])['timestamp'] / 1000000\n\n    token = record['token']\n    name = record['name']\n    date = datetime.utcfromtimestamp(start_time)\n    host = \"-\".join(record['name'].split(\"-\")[:2])\n    first_sample_token = record[\"first_sample_token\"]\n\n    entries.append((host, name, date, token, first_sample_token))\n            \ndf = pd.DataFrame(entries, columns=[\"host\", \"scene_name\", \"date\", \"scene_token\", \"first_sample_token\"])","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:43:59.956692Z","iopub.execute_input":"2024-09-10T13:43:59.95713Z","iopub.status.idle":"2024-09-10T13:43:59.97662Z","shell.execute_reply.started":"2024-09-10T13:43:59.957089Z","shell.execute_reply":"2024-09-10T13:43:59.975536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"host_count_df = df.groupby(\"host\")['scene_token'].count()\nprint(host_count_df)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:04.033662Z","iopub.execute_input":"2024-09-10T13:44:04.034393Z","iopub.status.idle":"2024-09-10T13:44:04.053981Z","shell.execute_reply.started":"2024-09-10T13:44:04.03435Z","shell.execute_reply":"2024-09-10T13:44:04.052384Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"validation_hosts = [\"host-a007\", \"host-a008\", \"host-a009\"]\n\nvalidation_df = df[df[\"host\"].isin(validation_hosts)]\nvi = validation_df.index\ntrain_df = df[~df.index.isin(vi)]\n\nprint(len(train_df), len(validation_df), \"train/validation split scene counts\")","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:07.710786Z","iopub.execute_input":"2024-09-10T13:44:07.711568Z","iopub.status.idle":"2024-09-10T13:44:07.724224Z","shell.execute_reply.started":"2024-09-10T13:44:07.711525Z","shell.execute_reply":"2024-09-10T13:44:07.723103Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample_token = train_df.first_sample_token.values[0]\nsample = level5data.get(\"sample\", sample_token)\n\nsample_lidar_token = sample[\"data\"][\"LIDAR_TOP\"]\nlidar_data = level5data.get(\"sample_data\", sample_lidar_token)\nlidar_filepath = level5data.get_sample_data_path(sample_lidar_token)\n\nego_pose = level5data.get(\"ego_pose\", lidar_data[\"ego_pose_token\"])\ncalibrated_sensor = level5data.get(\"calibrated_sensor\", lidar_data[\"calibrated_sensor_token\"])\n\n# Homogeneous transformation matrix from car frame to world frame.\nglobal_from_car = transform_matrix(ego_pose['translation'],\n                                   Quaternion(ego_pose['rotation']), inverse=False)\n\n# Homogeneous transformation matrix from sensor coordinate frame to ego car frame.\ncar_from_sensor = transform_matrix(calibrated_sensor['translation'], Quaternion(calibrated_sensor['rotation']),\n                                    inverse=False)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:09.62945Z","iopub.execute_input":"2024-09-10T13:44:09.630173Z","iopub.status.idle":"2024-09-10T13:44:09.640503Z","shell.execute_reply.started":"2024-09-10T13:44:09.630131Z","shell.execute_reply":"2024-09-10T13:44:09.639287Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"lidar_pointcloud = LidarPointCloud.from_file(lidar_filepath)\n\n# The lidar pointcloud is defined in the sensor's reference frame.\n# We want it in the car's reference frame, so we transform each point\nlidar_pointcloud.transform(car_from_sensor)\n\n# A sanity check, the points should be centered around 0 in car space.\nplt.hist(lidar_pointcloud.points[0], alpha=0.5, bins=30, label=\"X\")\nplt.hist(lidar_pointcloud.points[1], alpha=0.5, bins=30, label=\"Y\")\nplt.legend()\nplt.xlabel(\"Distance from car along axis\")\nplt.ylabel(\"Amount of points\")\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:13.5069Z","iopub.execute_input":"2024-09-10T13:44:13.507321Z","iopub.status.idle":"2024-09-10T13:44:14.053903Z","shell.execute_reply.started":"2024-09-10T13:44:13.507282Z","shell.execute_reply":"2024-09-10T13:44:14.052379Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def create_transformation_matrix_to_voxel_space(shape, voxel_size, offset):\n    \"\"\"\n    Constructs a transformation matrix given an output voxel shape such that (0,0,0) ends up in the center.\n    Voxel_size defines how large every voxel is in world coordinate, (1,1,1) would be the same as Minecraft voxels.\n    \n    An offset per axis in world coordinates (metric) can be provided, this is useful for Z (up-down) in lidar points.\n    \"\"\"\n    \n    shape, voxel_size, offset = np.array(shape), np.array(voxel_size), np.array(offset)\n    \n    tm = np.eye(4, dtype=np.float32)\n    translation = shape/2 + offset/voxel_size\n    \n    tm = tm * np.array(np.hstack((1/voxel_size, [1])))\n    tm[:3, 3] = np.transpose(translation)\n    return tm\n\ndef transform_points(points, transf_matrix):\n    \"\"\"\n    Transform (3,N) or (4,N) points using transformation matrix.\n    \"\"\"\n    if points.shape[0] not in [3,4]:\n        raise Exception(\"Points input should be (3,N) or (4,N) shape, received {}\".format(points.shape))\n    return transf_matrix.dot(np.vstack((points[:3, :], np.ones(points.shape[1]))))[:3, :]\n\n# Let's try it with some example values\ntm = create_transformation_matrix_to_voxel_space(shape=(100,100,4), voxel_size=(0.5,0.5,0.5), offset=(0,0,0.5))\np = transform_points(np.array([[10, 10, 0, 0, 0], [10, 5, 0, 0, 0],[0, 0, 0, 2, 0]], dtype=np.float32), tm)\nprint(p)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:15.616897Z","iopub.execute_input":"2024-09-10T13:44:15.617286Z","iopub.status.idle":"2024-09-10T13:44:15.630758Z","shell.execute_reply.started":"2024-09-10T13:44:15.617249Z","shell.execute_reply":"2024-09-10T13:44:15.629685Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def car_to_voxel_coords(points, shape, voxel_size, z_offset=0):\n    if len(shape) != 3:\n        raise Exception(\"Voxel volume shape should be 3 dimensions (x,y,z)\")\n        \n    if len(points.shape) != 2 or points.shape[0] not in [3, 4]:\n        raise Exception(\"Input points should be (3,N) or (4,N) in shape, found {}\".format(points.shape))\n\n    tm = create_transformation_matrix_to_voxel_space(shape, voxel_size, (0, 0, z_offset))\n    p = transform_points(points, tm)\n    return p\n\ndef create_voxel_pointcloud(points, shape, voxel_size=(0.5,0.5,1), z_offset=0):\n\n    points_voxel_coords = car_to_voxel_coords(points.copy(), shape, voxel_size, z_offset)\n    points_voxel_coords = points_voxel_coords[:3].transpose(1,0)\n    points_voxel_coords = np.int0(points_voxel_coords)\n    \n    bev = np.zeros(shape, dtype=np.float32)\n    bev_shape = np.array(shape)\n\n    within_bounds = (np.all(points_voxel_coords >= 0, axis=1) * np.all(points_voxel_coords < bev_shape, axis=1))\n    \n    points_voxel_coords = points_voxel_coords[within_bounds]\n    coord, count = np.unique(points_voxel_coords, axis=0, return_counts=True)\n        \n    # Note X and Y are flipped:\n    bev[coord[:,1], coord[:,0], coord[:,2]] = count\n    \n    return bev\n\ndef normalize_voxel_intensities(bev, max_intensity=16):\n    return (bev/max_intensity).clip(0,1)\n\n\nvoxel_size = (0.4,0.4,1.5)\nz_offset = -2.0\nbev_shape = (336, 336, 3)\n\nbev = create_voxel_pointcloud(lidar_pointcloud.points, bev_shape, voxel_size=voxel_size, z_offset=z_offset)\n\n# So that the values in the voxels range from 0,1 we set a maximum intensity.\nbev = normalize_voxel_intensities(bev)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:17.181906Z","iopub.execute_input":"2024-09-10T13:44:17.182621Z","iopub.status.idle":"2024-09-10T13:44:17.297413Z","shell.execute_reply.started":"2024-09-10T13:44:17.182579Z","shell.execute_reply":"2024-09-10T13:44:17.296362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.figure(figsize=(16,8))\nplt.imshow(bev)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:26.577631Z","iopub.execute_input":"2024-09-10T13:44:26.578564Z","iopub.status.idle":"2024-09-10T13:44:26.962903Z","shell.execute_reply.started":"2024-09-10T13:44:26.578519Z","shell.execute_reply":"2024-09-10T13:44:26.961921Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"boxes = level5data.get_boxes(sample_lidar_token)\n\ntarget_im = np.zeros(bev.shape[:3], dtype=np.uint8)\n\ndef move_boxes_to_car_space(boxes, ego_pose):\n    \"\"\"\n    Move boxes from world space to car space.\n    Note: mutates input boxes.\n    \"\"\"\n    translation = -np.array(ego_pose['translation'])\n    rotation = Quaternion(ego_pose['rotation']).inverse\n    \n    for box in boxes:\n        # Bring box to car space\n        box.translate(translation)\n        box.rotate(rotation)\n        \ndef scale_boxes(boxes, factor):\n    \"\"\"\n    Note: mutates input boxes\n    \"\"\"\n    for box in boxes:\n        box.wlh = box.wlh * factor\n\ndef draw_boxes(im, voxel_size, boxes, classes, z_offset=0.0):\n    for box in boxes:\n        # We only care about the bottom corners\n        corners = box.bottom_corners()\n        corners_voxel = car_to_voxel_coords(corners, im.shape, voxel_size, z_offset).transpose(1,0)\n        corners_voxel = corners_voxel[:,:2] # Drop z coord\n\n        class_color = classes.index(box.name) + 1\n        \n        if class_color == 0:\n            raise Exception(\"Unknown class: {}\".format(box.name))\n\n        cv2.drawContours(im, np.int0([corners_voxel]), 0, (class_color, class_color, class_color), -1)\n\n\n\nmove_boxes_to_car_space(boxes, ego_pose)\nscale_boxes(boxes, 0.8)\ndraw_boxes(target_im, voxel_size, boxes, classes, z_offset=z_offset)\n\nplt.figure(figsize=(8,8))\nplt.imshow((target_im > 0).astype(np.float32), cmap='Set2')\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:31.802233Z","iopub.execute_input":"2024-09-10T13:44:31.802624Z","iopub.status.idle":"2024-09-10T13:44:32.149447Z","shell.execute_reply.started":"2024-09-10T13:44:31.802586Z","shell.execute_reply":"2024-09-10T13:44:32.148287Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def visualize_lidar_of_sample(sample_token, axes_limit=80):\n    sample = level5data.get(\"sample\", sample_token)\n    sample_lidar_token = sample[\"data\"][\"LIDAR_TOP\"]\n    level5data.render_sample_data(sample_lidar_token, axes_limit=axes_limit)\n    \n# Don't worry about it being mirrored.\nvisualize_lidar_of_sample(sample_token)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:35.322756Z","iopub.execute_input":"2024-09-10T13:44:35.323164Z","iopub.status.idle":"2024-09-10T13:44:37.459265Z","shell.execute_reply.started":"2024-09-10T13:44:35.323126Z","shell.execute_reply":"2024-09-10T13:44:37.458243Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del bev, lidar_pointcloud, boxes","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:39.930801Z","iopub.execute_input":"2024-09-10T13:44:39.931713Z","iopub.status.idle":"2024-09-10T13:44:39.936078Z","shell.execute_reply.started":"2024-09-10T13:44:39.931671Z","shell.execute_reply":"2024-09-10T13:44:39.935081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Some hyperparameters we'll need to define for the system\nvoxel_size = (0.4, 0.4, 1.5)\nz_offset = -2.0\nbev_shape = (336, 336, 3)\n\n# We scale down each box so they are more separated when projected into our coarse voxel space.\nbox_scale = 0.8","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:41.189506Z","iopub.execute_input":"2024-09-10T13:44:41.190026Z","iopub.status.idle":"2024-09-10T13:44:41.195711Z","shell.execute_reply.started":"2024-09-10T13:44:41.189971Z","shell.execute_reply":"2024-09-10T13:44:41.194518Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# \"bev\" stands for birds eye view\ntrain_data_folder = os.path.join(ARTIFACTS_FOLDER, \"bev_train_data\")\nvalidation_data_folder = os.path.join(ARTIFACTS_FOLDER, \"./bev_validation_data\")","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:43.80599Z","iopub.execute_input":"2024-09-10T13:44:43.806402Z","iopub.status.idle":"2024-09-10T13:44:43.811477Z","shell.execute_reply.started":"2024-09-10T13:44:43.806364Z","shell.execute_reply":"2024-09-10T13:44:43.810378Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"NUM_WORKERS = os.cpu_count() * 3\n\ndef prepare_training_data_for_scene(first_sample_token, output_folder, bev_shape, voxel_size, z_offset, box_scale):\n    \"\"\"\n    Given a first sample token (in a scene), output rasterized input volumes and targets in birds-eye-view perspective.\n\n    \"\"\"\n    sample_token = first_sample_token\n    \n    while sample_token:\n        \n        sample = level5data.get(\"sample\", sample_token)\n\n        sample_lidar_token = sample[\"data\"][\"LIDAR_TOP\"]\n        lidar_data = level5data.get(\"sample_data\", sample_lidar_token)\n        lidar_filepath = level5data.get_sample_data_path(sample_lidar_token)\n\n        ego_pose = level5data.get(\"ego_pose\", lidar_data[\"ego_pose_token\"])\n        calibrated_sensor = level5data.get(\"calibrated_sensor\", lidar_data[\"calibrated_sensor_token\"])\n\n\n        global_from_car = transform_matrix(ego_pose['translation'],\n                                           Quaternion(ego_pose['rotation']), inverse=False)\n\n        car_from_sensor = transform_matrix(calibrated_sensor['translation'], Quaternion(calibrated_sensor['rotation']),\n                                            inverse=False)\n\n        try:\n            lidar_pointcloud = LidarPointCloud.from_file(lidar_filepath)\n            lidar_pointcloud.transform(car_from_sensor)\n        except Exception as e:\n            print (\"Failed to load Lidar Pointcloud for {}: {}:\".format(sample_token, e))\n            sample_token = sample[\"next\"]\n            continue\n        \n        bev = create_voxel_pointcloud(lidar_pointcloud.points, bev_shape, voxel_size=voxel_size, z_offset=z_offset)\n        bev = normalize_voxel_intensities(bev)\n\n        \n        boxes = level5data.get_boxes(sample_lidar_token)\n\n        target = np.zeros_like(bev)\n\n        move_boxes_to_car_space(boxes, ego_pose)\n        scale_boxes(boxes, box_scale)\n        draw_boxes(target, voxel_size, boxes=boxes, classes=classes, z_offset=z_offset)\n\n        bev_im = np.round(bev*255).astype(np.uint8)\n        target_im = target[:,:,0] # take one channel only\n\n        cv2.imwrite(os.path.join(output_folder, \"{}_input.png\".format(sample_token)), bev_im)\n        cv2.imwrite(os.path.join(output_folder, \"{}_target.png\".format(sample_token)), target_im)\n        \n        sample_token = sample[\"next\"]\n\nfor df, data_folder in [(train_df, train_data_folder), (validation_df, validation_data_folder)]:\n    print(\"Preparing data into {} using {} workers\".format(data_folder, NUM_WORKERS))\n    first_samples = df.first_sample_token.values\n\n    os.makedirs(data_folder, exist_ok=True)\n    \n    process_func = partial(prepare_training_data_for_scene,\n                           output_folder=data_folder, bev_shape=bev_shape, voxel_size=voxel_size, z_offset=z_offset, box_scale=box_scale)\n\n    pool = Pool(NUM_WORKERS)\n    for _ in tqdm_notebook(pool.imap_unordered(process_func, first_samples), total=len(first_samples)):\n        pass\n    pool.close()\n    del pool","metadata":{"execution":{"iopub.status.busy":"2024-09-10T13:44:44.959396Z","iopub.execute_input":"2024-09-10T13:44:44.960174Z","iopub.status.idle":"2024-09-10T14:18:52.417018Z","shell.execute_reply.started":"2024-09-10T13:44:44.960129Z","shell.execute_reply":"2024-09-10T14:18:52.415894Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.data\n\nclass BEVImageDataset(torch.utils.data.Dataset):\n    def __init__(self, input_filepaths, target_filepaths, map_filepaths=None):\n        self.input_filepaths = input_filepaths\n        self.target_filepaths = target_filepaths\n        self.map_filepaths = map_filepaths\n        \n        if map_filepaths is not None:\n            assert len(input_filepaths) == len(map_filepaths)\n        \n        assert len(input_filepaths) == len(target_filepaths)\n\n    def __len__(self):\n        return len(self.input_filepaths)\n\n    def __getitem__(self, idx):\n        input_filepath = self.input_filepaths[idx]\n        target_filepath = self.target_filepaths[idx]\n        \n        sample_token = input_filepath.split(\"/\")[-1].replace(\"_input.png\",\"\")\n        \n        im = cv2.imread(input_filepath, cv2.IMREAD_UNCHANGED)\n        \n        if self.map_filepaths:\n            map_filepath = self.map_filepaths[idx]\n            map_im = cv2.imread(map_filepath, cv2.IMREAD_UNCHANGED)\n            im = np.concatenate((im, map_im), axis=2)\n        \n        target = cv2.imread(target_filepath, cv2.IMREAD_UNCHANGED)\n        \n        im = im.astype(np.float32)/255\n        target = target.astype(np.int64)\n        \n        im = torch.from_numpy(im.transpose(2,0,1))\n        target = torch.from_numpy(target)\n        \n        return im, target, sample_token\n\ninput_filepaths = sorted(glob.glob(os.path.join(train_data_folder, \"*_input.png\")))\ntarget_filepaths = sorted(glob.glob(os.path.join(train_data_folder, \"*_target.png\")))\n\ntrain_dataset = BEVImageDataset(input_filepaths, target_filepaths)\n    \nim, target, sample_token = train_dataset[1]\nim = im.numpy()\ntarget = target.numpy()\n\nplt.figure(figsize=(16,8))\n\ntarget_as_rgb = np.repeat(target[...,None], 3, 2)\n# Transpose the input volume CXY to XYC order, which is what matplotlib requires.\nplt.imshow(np.hstack((im.transpose(1,2,0)[...,:3], target_as_rgb)))\nplt.title(sample_token)\nplt.show()\n\nvisualize_lidar_of_sample(sample_token)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T14:19:26.513455Z","iopub.execute_input":"2024-09-10T14:19:26.513894Z","iopub.status.idle":"2024-09-10T14:19:33.907175Z","shell.execute_reply.started":"2024-09-10T14:19:26.513845Z","shell.execute_reply":"2024-09-10T14:19:33.906065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class UNet(nn.Module):\n    def __init__(\n        self,\n        in_channels=1,#number of input channels\n        n_classes=2,#number of output channels\n        depth=5, #depth of the network\n        wf=6, #number of filters in the first layer is 2**wf\n        padding=False,#if True, apply padding such that the input shape is the same as the output.\n        batch_norm=False, # Use Batch Normlization after layers with an activation function\n        up_mode='upconv', # one of 'upconv' or 'upsample'.\n                            #'upconv' will use transposed convolutions for learned upsampling.\n                            #'upsample' will use bilinear upsampling.\n    ):\n        super(UNet, self).__init__()\n        assert up_mode in ('upconv', 'upsample')\n        self.padding = padding\n        self.depth = depth\n        prev_ch = in_channels\n        self.down_path = nn.ModuleList()\n        for i in range(depth):\n            self.down_path.append(\n                UNetConvBlock(prev_ch, 2 ** (wf + i), padding, batch_norm)\n            )\n            prev_ch = 2 ** (wf + i)\n\n        self.up_path = nn.ModuleList()\n        for i in reversed(range(depth - 1)):\n            self.up_path.append(\n                UNetUpBlock(prev_ch, 2 ** (wf + i), up_mode, padding, batch_norm)\n            )\n            prev_ch = 2 ** (wf + i)\n\n        self.last = nn.Conv2d(prev_ch, n_classes, kernel_size=1)\n\n    def forward(self, x):\n        blocks = []\n        for i, down in enumerate(self.down_path):\n            x = down(x)\n            if i != len(self.down_path) - 1:\n                blocks.append(x)\n                x = F.max_pool2d(x, 2)\n\n        for i, up in enumerate(self.up_path):\n            x = up(x, blocks[-i - 1])\n\n        return self.last(x)\n\n\nclass UNetConvBlock(nn.Module):\n    \"\"\"\n    class defining each UNet Convolution block\n    \n    Consists of 2 sets 2-D convolution layer followed by a ReLU activation layer \n    and a batch normalization layer, if batch_norm flag is true, in sequential model \n    \n    \"\"\"\n    def __init__(self, in_size, out_size, padding, batch_norm):\n        super(UNetConvBlock, self).__init__()\n        convblock = []\n\n        convblock.append(nn.Conv2d(in_size, out_size, kernel_size=3, padding=int(padding)))\n        convblock.append(nn.ReLU())\n        if batch_norm:\n            convblock.append(nn.BatchNorm2d(out_size))\n\n        convblock.append(nn.Conv2d(out_size, out_size, kernel_size=3, padding=int(padding)))\n        convblock.append(nn.ReLU())\n        if batch_norm:\n            convblock.append(nn.BatchNorm2d(out_size))\n\n        self.convblock = nn.Sequential(*convblock)\n\n    def forward(self, x):\n        out = self.convblock(x)\n        return out\n\n\nclass UNetUpBlock(nn.Module):\n    \"\"\"\n        Based on up_mode the UNet up block is defined.\n        upconv : 2-D Convolution transpose layer\n        upsample : sequentially arranged upsampling layer and 2D convolution layer\n    \n    \"\"\"\n    def __init__(self, in_size, out_size, up_mode, padding, batch_norm):\n        super(UNetUpBlock, self).__init__()\n        if up_mode == 'upconv':\n            self.up = nn.ConvTranspose2d(in_size, out_size, kernel_size=2, stride=2)\n        elif up_mode == 'upsample':\n            self.up = nn.Sequential(\n                nn.Upsample(mode='bilinear', scale_factor=2),\n                nn.Conv2d(in_size, out_size, kernel_size=1),\n            )\n\n        self.conv_block = UNetConvBlock(in_size, out_size, padding, batch_norm)\n\n    def center_crop(self, layer, target_size):\n        _, _, layer_height, layer_width = layer.size()\n        diff_y = (layer_height - target_size[0]) // 2\n        diff_x = (layer_width - target_size[1]) // 2\n        return layer[\n            :, :, diff_y : (diff_y + target_size[0]), diff_x : (diff_x + target_size[1])\n        ]\n\n    def forward(self, x, bridge):\n        up = self.up(x)\n        crop1 = self.center_crop(bridge, up.shape[2:])\n        out = torch.cat([up, crop1], 1)\n        out = self.conv_block(out)\n\n        return out","metadata":{"execution":{"iopub.status.busy":"2024-09-10T14:19:49.889111Z","iopub.execute_input":"2024-09-10T14:19:49.889718Z","iopub.status.idle":"2024-09-10T14:19:49.912829Z","shell.execute_reply.started":"2024-09-10T14:19:49.889677Z","shell.execute_reply":"2024-09-10T14:19:49.911624Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_unet_model(in_channels=3, num_output_classes=2):\n    model = UNet(in_channels=in_channels, n_classes=num_output_classes, wf=5, depth=4, padding=True, up_mode='upsample')\n    \n    # Optional, for multi GPU training and inference\n    model = nn.DataParallel(model)\n    return model","metadata":{"execution":{"iopub.status.busy":"2024-09-10T14:20:00.908162Z","iopub.execute_input":"2024-09-10T14:20:00.908582Z","iopub.status.idle":"2024-09-10T14:20:00.914537Z","shell.execute_reply.started":"2024-09-10T14:20:00.908543Z","shell.execute_reply":"2024-09-10T14:20:00.913434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def visualize_predictions(input_image, prediction, target, n_images=2, apply_softmax=True):\n    \"\"\"\n    Takes as input 3 PyTorch tensors, plots the input image, predictions and targets.\n    \"\"\"\n    # Only select the first n images\n    prediction = prediction[:n_images]\n    target = target[:n_images]\n    input_image = input_image[:n_images]\n\n    prediction = prediction.detach().cpu().numpy()\n    if apply_softmax:\n        prediction = scipy.special.softmax(prediction, axis=1)\n    class_one_preds = np.hstack(1-prediction[:,0])\n\n    target = np.hstack(target.detach().cpu().numpy())\n\n    class_rgb = np.repeat(class_one_preds[..., None], 3, axis=2)\n    class_rgb[...,2] = 0\n    class_rgb[...,1] = target\n\n    \n    input_im = np.hstack(input_image.cpu().numpy().transpose(0,2,3,1))\n    \n    if input_im.shape[2] == 3:\n        input_im_grayscale = np.repeat(input_im.mean(axis=2)[..., None], 3, axis=2)\n        overlayed_im = (input_im_grayscale*0.6 + class_rgb*0.7).clip(0,1)\n    else:\n        input_map = input_im[...,3:]\n        overlayed_im = (input_map*0.6 + class_rgb*0.7).clip(0,1)\n\n    thresholded_pred = np.repeat(class_one_preds[..., None] > 0.5, 3, axis=2)\n\n    fig = plt.figure(figsize=(12,26))\n    plot_im = np.vstack([class_rgb, input_im[...,:3], overlayed_im, thresholded_pred]).clip(0,1).astype(np.float32)\n    plt.imshow(plot_im)\n    plt.axis(\"off\")\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2024-09-10T14:20:02.085921Z","iopub.execute_input":"2024-09-10T14:20:02.086964Z","iopub.status.idle":"2024-09-10T14:20:02.098938Z","shell.execute_reply.started":"2024-09-10T14:20:02.08692Z","shell.execute_reply":"2024-09-10T14:20:02.097924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# We weigh the loss for the 0 class lower to account for (some of) the big class imbalance.\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nclass_weights = torch.from_numpy(np.array([0.2] + [1.0]*len(classes), dtype=np.float32))\nclass_weights = class_weights.to(device)","metadata":{"execution":{"iopub.status.busy":"2024-09-10T14:20:03.885502Z","iopub.execute_input":"2024-09-10T14:20:03.886282Z","iopub.status.idle":"2024-09-10T14:20:04.185993Z","shell.execute_reply.started":"2024-09-10T14:20:03.886234Z","shell.execute_reply":"2024-09-10T14:20:04.184717Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"batch_size = 8\nepochs = 15 # Note: We may be able to train for longer and expect better results, the reason this number is low is to keep the runtime short.\n\nmodel = get_unet_model(num_output_classes=len(classes)+1)\nmodel = model.to(device)\n\noptim = torch.optim.Adam(model.parameters(), lr=1e-3)\ndataloader = torch.utils.data.DataLoader(train_dataset, batch_size, shuffle=True, num_workers=os.cpu_count()*2)\n\nall_losses = []\n\nfor epoch in range(1, epochs+1):\n    print(\"Epoch\", epoch)\n    \n    epoch_losses = []\n    progress_bar = tqdm_notebook(dataloader)\n    \n    for ii, (X, target, sample_ids) in enumerate(progress_bar):\n        X = X.to(device)  # [N, 3, H, W]\n        target = target.to(device)  # [N, H, W] with class indices (0, 1)\n        prediction = model(X)  # [N, 2, H, W]\n        loss = F.cross_entropy(prediction, target, weight=class_weights)\n\n        optim.zero_grad()\n        loss.backward()\n        optim.step()\n        \n        epoch_losses.append(loss.detach().cpu().numpy())\n\n        if ii == 0:\n            visualize_predictions(X, prediction, target)\n    \n    print(\"Loss:\", np.mean(epoch_losses))\n    all_losses.extend(epoch_losses)\n    \n    checkpoint_filename = \"unet_checkpoint_epoch_{}.pth\".format(epoch)\n    checkpoint_filepath = os.path.join(ARTIFACTS_FOLDER, checkpoint_filename)\n    torch.save(model.state_dict(), checkpoint_filepath)\n    \nplt.figure(figsize=(12,12))\nplt.plot(all_losses, alpha=0.75)\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-09-10T14:20:05.212158Z","iopub.execute_input":"2024-09-10T14:20:05.212571Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}