{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Motion Prediction with PointNet"},{"metadata":{},"cell_type":"markdown","source":"# About"},{"metadata":{},"cell_type":"markdown","source":"In this work, I will be using the [pointnet architecture](https://github.com/charlesq34/pointnet) to predict autonomous vehicle motions. This work is completely different from the main public kernels and exposes another way of dealing with motion prediction.\n\nOne of the edges of the pointnet architecture is that it works directly on the agents set and, therefore, no rastarization is required. Hence, we can can get a decent accuracy  while training only on the small train set and just for 5 to 7 hours. Hence, this model could be run even if one doesn't have huge GPU capacities. But, if you do have enough GPU, you can increase the model size and train on the full train set."},{"metadata":{},"cell_type":"markdown","source":"### This kernel is inference only, the training process is heavier and is on its road !"},{"metadata":{},"cell_type":"markdown","source":"The PointNet model uses a combination of feed forward (conv1d) models along with somme transformation matrices and symmetric functions (max pooling) to deal with unordered set of points (agents). For more reading, please visit the [PointNet project page](https://stanford.edu/~rqi/pointnet/) or [the original paper](https://arxiv.org/pdf/1612.00593.pdf). My code is inspired from [this github repos](https://github.com/charlesq34/pointnet)."},{"metadata":{},"cell_type":"markdown","source":"![](https://stanford.edu/~rqi/pointnet/images/pointnet.jpg)"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"> PointNet can be applied for classification and image segmentation. Since we will be needing an output for each point (agent), only the segmentation version will be suitable.\n\n\n![PointNet architecture](https://github.com/charlesq34/pointnet/blob/master/doc/teaser.png?raw=true)"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"!pip install zarr > /dev/null\n!pip install  --use-feature=2020-resolver pytorch-lightning  > /dev/null","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import zarr\nfrom abc import ABC\nfrom pathlib import Path\nfrom numcodecs import blosc\nimport pandas as pd, numpy as np\n\nimport bisect\nimport itertools as it\nfrom tqdm.notebook import tqdm\n\n\nimport torch\nfrom torch import nn, optim \nimport torch.nn.functional as F\nfrom torch.autograd import Variable\nfrom torch.utils.data import Dataset, DataLoader\n\nimport pickle, copy, re,time, datetime, random, warnings\n\npd.options.display.max_columns=305","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# blosc.set_nthreads(6)  \n# blosc.use_threads = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DATA_ROOT = Path(\"../input/lyft-motion-prediction-autonomous-vehicles\")\nZARR_PATH = Path(\"scenes/test.zarr\")\nprint(\"DATA_ROOT: {}\\nZARR_PATH: {}\".format(DATA_ROOT, ZARR_PATH))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"HBACKWARD = 15\nHFORWARD = 0\nNFRAMES = 1\nFRAME_STRIDE = 0\nAGENT_FEATURE_DIM = 8\nMAX_AGENTS = 150","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"TIME_FORMAT = r\"%Y-%m-%dT%H:%M:%S%Z\"\ndef get_utc():\n    return datetime.datetime.now(datetime.timezone.utc).strftime(TIME_FORMAT)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Dataset"},{"metadata":{"trusted":true},"cell_type":"code","source":"PERCEPTION_LABELS = [\n    \"PERCEPTION_LABEL_NOT_SET\",\n    \"PERCEPTION_LABEL_UNKNOWN\",\n    \"PERCEPTION_LABEL_DONTCARE\",\n    \"PERCEPTION_LABEL_CAR\",\n    \"PERCEPTION_LABEL_VAN\",\n    \"PERCEPTION_LABEL_TRAM\",\n    \"PERCEPTION_LABEL_BUS\",\n    \"PERCEPTION_LABEL_TRUCK\",\n    \"PERCEPTION_LABEL_EMERGENCY_VEHICLE\",\n    \"PERCEPTION_LABEL_OTHER_VEHICLE\",\n    \"PERCEPTION_LABEL_BICYCLE\",\n    \"PERCEPTION_LABEL_MOTORCYCLE\",\n    \"PERCEPTION_LABEL_CYCLIST\",\n    \"PERCEPTION_LABEL_MOTORCYCLIST\",\n    \"PERCEPTION_LABEL_PEDESTRIAN\",\n    \"PERCEPTION_LABEL_ANIMAL\",\n    \"AVRESEARCH_LABEL_DONTCARE\",\n]\nKEPT_PERCEPTION_LABELS = [\n    \"PERCEPTION_LABEL_UNKNOWN\",\n    \"PERCEPTION_LABEL_CAR\",\n    \"PERCEPTION_LABEL_CYCLIST\",\n    \"PERCEPTION_LABEL_PEDESTRIAN\",\n]\nKEPT_PERCEPTION_LABELS_DICT = {label:PERCEPTION_LABELS.index(label) for label in KEPT_PERCEPTION_LABELS}\nKEPT_PERCEPTION_KEYS = sorted(KEPT_PERCEPTION_LABELS_DICT.values())","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LabelEncoder:\n    def  __init__(self, max_size=500, default_val=-1):\n        self.max_size = max_size\n        self.labels = {}\n        self.default_val = default_val\n\n    @property\n    def nlabels(self):\n        return len(self.labels)\n\n    def reset(self):\n        self.labels = {}\n\n    def partial_fit(self, keys):\n        nlabels = self.nlabels\n        available = self.max_size - nlabels\n\n        if available < 1:\n            return\n\n        keys = set(keys)\n        new_keys = list(keys - set(self.labels))\n\n        if not len(new_keys):\n            return\n        \n        self.labels.update(dict(zip(new_keys, range(nlabels, nlabels + available) )))\n    \n    def fit(self, keys):\n        self.reset()\n        self.partial_fit(keys)\n\n    def get(self, key):\n        return self.labels.get(key, self.default_val)\n    \n    def transform(self, keys):\n        return np.array(list(map(self.get, keys)))\n\n    def fit_transform(self, keys, partial=True):\n        self.partial_fit(keys) if partial else self.fit(keys)\n        return self.transform(keys)","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"class CustomLyftDataset(Dataset):\n    feature_mins = np.array([-17.336, -27.137, 0. , 0., 0. , -3.142, -37.833, -65.583],\n    dtype=\"float32\")[None,None, None]\n\n    feature_maxs = np.array([17.114, 20.787, 42.854, 42.138,  7.079,  3.142, 29.802, 35.722],\n    dtype=\"float32\")[None,None, None]\n\n\n\n    def __init__(self, zdataset, scenes=None, nframes=10, frame_stride=15, hbackward=10, \n                 hforward=50, max_agents=150, agent_feature_dim=8):\n        \"\"\"\n        Custom Lyft dataset reader.\n        \n        Parmeters:\n        ----------\n        zdataset: zarr dataset\n            The root dataset, containing scenes, frames and agents\n            \n        nframes: int\n            Number of frames per scene\n            \n        frame_stride: int\n            The stride when reading the **nframes** frames from a scene\n            \n        hbackward: int\n            Number of backward frames from  current frame\n            \n        hforward: int\n            Number forward frames from current frame\n        \n        max_agents: int \n            Max number of agents to read for each target frame. Note that,\n            this also include the backward agents but not the forward ones.\n        \"\"\"\n        super().__init__()\n        self.zdataset = zdataset\n        self.scenes = scenes if scenes is not None else []\n        self.nframes = nframes\n        self.frame_stride = frame_stride\n        self.hbackward = hbackward\n        self.hforward = hforward\n        self.max_agents = max_agents\n\n        self.nread_frames = (nframes-1)*frame_stride + hbackward + hforward\n\n        self.frame_fields = ['timestamp', 'agent_index_interval']\n\n        self.agent_feature_dim = agent_feature_dim\n\n        self.filter_scenes()\n      \n    def __len__(self):\n        return len(self.scenes)\n\n    def filter_scenes(self):\n        self.scenes = [scene for scene in self.scenes if self.get_nframes(scene) > self.nread_frames]\n\n\n    def __getitem__(self, index):\n        return self.read_frames(scene=self.scenes[index])\n\n    def get_nframes(self, scene, start=None):\n        frame_start = scene[\"frame_index_interval\"][0]\n        frame_end = scene[\"frame_index_interval\"][1]\n        nframes = (frame_end - frame_start) if start is None else ( frame_end - max(frame_start, start) )\n        return nframes\n\n\n    def _read_frames(self, scene, start=None):\n        nframes = self.get_nframes(scene, start=start)\n        assert nframes >= self.nread_frames\n\n        frame_start = scene[\"frame_index_interval\"][0]\n\n        start = start or frame_start + np.random.choice(nframes-self.nread_frames)\n        frames = self.zdataset.frames.get_basic_selection(\n            selection=slice(start, start+self.nread_frames),\n            fields=self.frame_fields,\n            )\n        return frames\n    \n\n    def parse_frame(self, frame):\n        return frame\n\n    def parse_agent(self, agent):\n        return agent\n\n    def read_frames(self, scene, start=None,  white_tracks=None, encoder=False):\n        white_tracks = white_tracks or []\n        frames = self._read_frames(scene=scene, start=start)\n\n        agent_start = frames[0][\"agent_index_interval\"][0]\n        agent_end = frames[-1][\"agent_index_interval\"][1]\n\n        agents = self.zdataset.agents[agent_start:agent_end]\n\n\n        X = np.zeros((self.nframes, self.max_agents, self.hbackward, self.agent_feature_dim), dtype=np.float32)\n        target = np.zeros((self.nframes, self.max_agents, self.hforward, 2),  dtype=np.float32)\n        target_availability = np.zeros((self.nframes, self.max_agents, self.hforward), dtype=np.uint8)\n        X_availability = np.zeros((self.nframes, self.max_agents, self.hbackward), dtype=np.uint8)\n\n        for f in range(self.nframes):\n            backward_frame_start = f*self.frame_stride\n            forward_frame_start = f*self.frame_stride+self.hbackward\n            backward_frames = frames[backward_frame_start:backward_frame_start+self.hbackward]\n            forward_frames = frames[forward_frame_start:forward_frame_start+self.hforward]\n\n            backward_agent_start = backward_frames[-1][\"agent_index_interval\"][0] - agent_start\n            backward_agent_end = backward_frames[-1][\"agent_index_interval\"][1] - agent_start\n\n            backward_agents = agents[backward_agent_start:backward_agent_end]\n\n            le = LabelEncoder(max_size=self.max_agents)\n            le.fit(white_tracks)\n            le.partial_fit(backward_agents[\"track_id\"])\n\n            for iframe, frame in enumerate(backward_frames):\n                backward_agent_start = frame[\"agent_index_interval\"][0] - agent_start\n                backward_agent_end = frame[\"agent_index_interval\"][1] - agent_start\n\n                backward_agents = agents[backward_agent_start:backward_agent_end]\n\n                track_ids = le.transform(backward_agents[\"track_id\"])\n                mask = (track_ids != le.default_val)\n                mask_agents = backward_agents[mask]\n                mask_ids = track_ids[mask]\n                X[f, mask_ids, iframe, :2] = mask_agents[\"centroid\"]\n                X[f, mask_ids, iframe, 2:5] = mask_agents[\"extent\"]\n                X[f, mask_ids, iframe, 5] = mask_agents[\"yaw\"]\n                X[f, mask_ids, iframe, 6:8] = mask_agents[\"velocity\"]\n\n                X_availability[f, mask_ids, iframe] = 1\n\n            \n            for iframe, frame in enumerate(forward_frames):\n                forward_agent_start = frame[\"agent_index_interval\"][0] - agent_start\n                forward_agent_end = frame[\"agent_index_interval\"][1] - agent_start\n\n                forward_agents = agents[forward_agent_start:forward_agent_end]\n\n                track_ids = le.transform(forward_agents[\"track_id\"])\n                mask = track_ids != le.default_val\n\n                target[f, track_ids[mask], iframe] = forward_agents[mask][\"centroid\"]\n                target_availability[f, track_ids[mask], iframe] = 1\n\n        target -= X[:,:,[-1], :2]\n        target *= target_availability[:,:,:,None]\n        X[:,:,:, :2] -= X[:,:,[-1], :2]\n        X *= X_availability[:,:,:,None]\n        X -= self.feature_mins\n        X /= (self.feature_maxs - self.feature_mins)\n\n        if encoder:\n            return X, target, target_availability, le\n        return X, target, target_availability","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":false},"cell_type":"markdown","source":"# PointNet"},{"metadata":{"trusted":true},"cell_type":"code","source":"class STNkd(nn.Module):\n    def __init__(self,  k=64):\n        super(STNkd, self).__init__()\n        \n        self.conv = nn.Sequential(\n            nn.Conv1d(k, 256, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(256, 256, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(256, 512, kernel_size=1), nn.ReLU(),\n        )\n        \n        self.fc = nn.Sequential(\n            nn.Linear(512, k*k),nn.ReLU(),\n        )\n        self.k = k\n\n    def forward(self, x):\n        batchsize = x.size()[0]\n        x = self.conv(x)\n        x = torch.max(x, 2)[0]\n        x = self.fc(x)\n\n        iden = Variable(torch.from_numpy(np.eye(self.k).flatten().astype(np.float32))).view(1,\n                                                                            self.k*self.k).repeat(batchsize,1)\n        if x.is_cuda:\n            iden = iden.cuda()\n        x = x + iden\n        x = x.view(-1, self.k, self.k)\n        return x","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class PointNetfeat(nn.Module):\n    def __init__(self, global_feat = False, feature_transform = False, stn1_dim = 120,\n                 stn2_dim = 64):\n        super(PointNetfeat, self).__init__()\n        \n        self.global_feat = global_feat\n        self.feature_transform = feature_transform\n        self.stn1_dim = stn1_dim\n        self.stn2_dim = stn2_dim\n        \n        self.stn = STNkd(k=stn1_dim)\n        \n        self.conv1 = nn.Sequential(\n            nn.Conv1d(stn1_dim, 256, kernel_size=1), nn.ReLU(),\n        )\n        \n        self.conv2 = nn.Sequential(\n            nn.Conv1d(256, 256, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(256, 1024, kernel_size=1), nn.ReLU(),\n            nn.Conv1d(1024, 2048, kernel_size=1), nn.ReLU(),\n        )\n        \n        if self.feature_transform:\n            self.fstn = STNkd(k=stn2_dim)\n\n    def forward(self, x):\n        n_pts = x.size()[2]\n        trans = self.stn(x)\n        x = x.transpose(2, 1)\n        x = torch.bmm(x, trans)\n        x = x.transpose(2, 1)\n        \n        x = self.conv1(x)\n\n        if self.feature_transform:\n            trans_feat = self.fstn(x)\n            x = x.transpose(2,1)\n            x = torch.bmm(x, trans_feat)\n            x = x.transpose(2,1)\n        else:\n            trans_feat = None\n\n        pointfeat = x\n        \n        x = self.conv2(x)\n        x = torch.max(x, 2)[0]\n        if self.global_feat:\n            return x, trans, trans_feat\n        else:\n            x = x[:,:,None].repeat(1, 1, n_pts)\n            return torch.cat([x, pointfeat], 1), trans, trans_feat","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class LyftNet(nn.Module):   \n    def __init__(self):\n        super().__init__()\n        \n        self.pnet = PointNetfeat()\n\n        self.fc0 = nn.Sequential(\n            nn.Linear(2048+256, 1024), nn.ReLU(),\n        )\n\n        self.fc = nn.Sequential(\n            nn.Linear(1024, 300),\n        )\n\n        self.c_net = nn.Sequential(\n            nn.Linear(1024, 3),\n        )\n        \n    \n    def forward(self, x):\n        bsize, npoints, hb, nf = x.shape \n        \n        # Push points to the last  dim\n        x = x.transpose(1, 3)\n\n        # Merge time with features\n        x = x.reshape(bsize, hb*nf, npoints)\n\n        x, trans, trans_feat = self.pnet(x)\n\n        # Push featuresxtime to the last dim\n        x = x.transpose(1,2)\n\n        x = self.fc0(x)\n\n        c = self.c_net(x)\n        x = self.fc(x)\n\n        return c,x","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Prediction"},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\nmask = np.load(DATA_ROOT/\"scenes/mask.npz\")['arr_0']\nagent_ids = np.where(mask)[0]\nagent_ids","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"z = zarr.open(DATA_ROOT.joinpath(ZARR_PATH).as_posix())\nframes_ij = z.scenes[\"frame_index_interval\"]\nagents_ij = z.frames[\"agent_index_interval\"]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_scene(agent_id):\n    frame_id = bisect.bisect_right(agents_ij[:, 0], agent_id)-1\n    scene_id = bisect.bisect_right(frames_ij[:, 0], frame_id)-1\n    \n    scene = z.scenes[scene_id]\n    frame = z.frames[frame_id]\n    agent = z.agents[agent_id]\n    return scene,(frame,frame_id),agent","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"agent_id = agent_ids[np.random.choice(len(agent_ids))]\nscene,(frame,frame_id), agent = get_scene(agent_id)\nscene,(frame,frame_id), agent[\"track_id\"]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dt = CustomLyftDataset(\n        z, \n        nframes=NFRAMES,\n        frame_stride=FRAME_STRIDE,\n        hbackward=HBACKWARD,\n        hforward=HFORWARD,\n        max_agents=MAX_AGENTS,\n        agent_feature_dim=AGENT_FEATURE_DIM,\n)\ndt.nread_frames","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"paths = [\n    \"../input/neural-net-on-lyft-tabular-data/lyfnet_epoch13.ckpt\",\n]\nweights = np.array([1.0])\nweights /= weights.sum()\nnets = []\nfor path in paths:\n    net = LyftNet().to(DEVICE)\n    net.load_state_dict(torch.load(path, map_location=DEVICE)[\"state_dict\"])\n    net.eval()\n    nets.append(net)\n\nnets[0]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def make_colnames():\n    cols= list(\n        it.chain(*[\n    [f\"coord_x0{i}\", f\"coord_y0{i}\", f\"coord_x1{i}\", f\"coord_y1{i}\", f\"coord_x2{i}\", f\"coord_y2{i}\"]\n    for i in range(50)]))\n    return [\"timestamp\", \"track_id\"] + [\"conf_0\", \"conf_1\", \"conf_2\"] + cols","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"@torch.no_grad()\ndef predict(agent_id):\n    scene, (frame,frame_id), agent = get_scene(agent_id)\n    c = np.zeros(3)\n    y = np.zeros(300)\n    try:\n        X, _,_, le = dt.read_frames(scene,\n            start=frame_id-HBACKWARD+1,\n            white_tracks=[agent[\"track_id\"]],\n            encoder=True\n        )\n        \n        X = torch.from_numpy(X).to(DEVICE)\n        cs,ys=  [],[]\n        for net in nets:\n            c, y = net(X)\n            c = torch.softmax(c[0, [le.labels[agent[\"track_id\"]]]],dim=1)[0].cpu().numpy()\n            y = y[0, le.labels[agent[\"track_id\"]]].cpu().numpy()\n            cs.append(c)\n            ys.append(y)\n            \n        cs = np.sum(np.vstack(cs)*weights[None], axis=1)\n        ys = np.vstack(ys).reshape(len(ys), 50, 3, 2)\n        ys = np.sum(ys*weights[None,None,:,None], axis=2).reshape(len(ys), -1)\n        \n    except Exception as e:\n        print(f'Exception : {e}')\n#         raise ValueError() from e\n    return [frame[\"timestamp\"]],[agent[\"track_id\"]],c,y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ntimestamp,track_id, c,y = [],[],[],[]\n\nfor icount, agent_id in tqdm(list(enumerate(agent_ids))):\n    _timestamp, _track_id, _c,_y = predict(agent_id)\n    timestamp.extend(_timestamp)\n    track_id.extend(_track_id)\n    c.append(_c)\n    y.append(_y)\nc = np.vstack(c)\ny = np.vstack(y)\ny.shape, c.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"%%time\n\ncols = make_colnames()\ndf = pd.DataFrame(columns=cols)\ndf[cols[0]] = timestamp\ndf[cols[1]] = track_id\ndf[cols[2:5]] = c\ndf[cols[5:]] = y\ndf.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df.head(20)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"df.to_csv(\"submission.csv\", index=False)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"<div style=\"text-align: center\"><strong>Thanks</strong></div>"},{"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}