From 043223c22f318eb9edc3b28de21a31c96c418cae Mon Sep 17 00:00:00 2001
From: Tristan Rice <rice@fn.lc>
Date: Mon, 23 Oct 2023 22:35:44 -0700
Subject: [PATCH] nuscenes: revert custom changes

---
 notebooks/dataset-debug.ipynb           | 48 +++++++++++++++++++++----
 torchdrive/datasets/nuscenes_dataset.py | 22 ++++++++----
 2 files changed, 58 insertions(+), 12 deletions(-)

diff --git a/notebooks/dataset-debug.ipynb b/notebooks/dataset-debug.ipynb
index 3c9eb3a..30cc109 100644
--- a/notebooks/dataset-debug.ipynb
+++ b/notebooks/dataset-debug.ipynb
@@ -6,8 +6,8 @@
    "id": "08b1582e",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2023-10-23T17:34:55.660315Z",
-     "start_time": "2023-10-23T17:34:55.600897Z"
+     "end_time": "2023-10-23T18:13:16.144669Z",
+     "start_time": "2023-10-23T18:13:16.121225Z"
     }
    },
    "outputs": [],
@@ -17,12 +17,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 35,
+   "execution_count": 19,
    "id": "32dfa27f",
    "metadata": {
     "ExecuteTime": {
-     "end_time": "2023-10-23T17:55:42.881190Z",
-     "start_time": "2023-10-23T17:55:33.366184Z"
+     "end_time": "2023-10-23T18:24:22.822325Z",
+     "start_time": "2023-10-23T18:24:19.700327Z"
     },
     "scrolled": true
    },
@@ -46,7 +46,7 @@
       "31206 sample_data,\n",
       "18538 sample_annotation,\n",
       "4 map,\n",
-      "Done loading in 0.478 seconds.\n",
+      "Done loading in 0.318 seconds.\n",
       "======\n",
       "Reverse indexing ...\n",
       "Found 2312 scenes for CAM_FRONT\n",
@@ -108,6 +108,42 @@
     "#    display_img(example.color[cam][0].float())"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "8ccf96b4",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2023-10-23T18:24:25.389609Z",
+     "start_time": "2023-10-23T18:24:25.366319Z"
+    }
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "keys\n",
+       "f4d2a6c281f34a7eb8bb033d82321f79      0\n",
+       "2ee327ac0903407dbb42c754861c1e63      1\n",
+       "341ad78c836f4dc4bd2c61ad12285a63      2\n",
+       "794da4b758054140bd84d26b1146e46e      3\n",
+       "eff1fe2c0eaa44958bf11f7398e8293b      4\n",
+       "                                   ... \n",
+       "302e35d6ff604f17adf2ccdebce473d1    115\n",
+       "6a46edcebe58409798fbfa062a05cbc3    116\n",
+       "d554c8b1480e42679bec69c7b29f3c2a    117\n",
+       "2dbc97f826174617b24cc75707c580b5    118\n",
+       "a858134f587548429f781b26b3348d1f    119\n",
+       "Name: values, Length: 120, dtype: int64"
+      ]
+     },
+     "execution_count": 20,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": []
+  },
   {
    "cell_type": "code",
    "execution_count": 3,
diff --git a/torchdrive/datasets/nuscenes_dataset.py b/torchdrive/datasets/nuscenes_dataset.py
index 817034c..baaddf7 100644
--- a/torchdrive/datasets/nuscenes_dataset.py
+++ b/torchdrive/datasets/nuscenes_dataset.py
@@ -49,8 +49,7 @@ def as_py(data: object) -> object:
     return data
 
 
-# NuScenes = UnpatchedNuScenes
-class NuScenes(UnpatchedNuScenes):
+class CustomNuScenes(UnpatchedNuScenes):
     def __init__(
         self,
         version: str = "v1.0-mini",
@@ -60,8 +59,12 @@ def __init__(
     ):
         super().__init__(version, dataroot, verbose, map_resolution)
 
+        to_delete = ["map", "sample_annotation", "instance"]
+        for attr in to_delete:
+            delattr(self, attr)
+
         for table in self.table_names:
-            if table == "map":
+            if table in to_delete:
                 continue
             setattr(self, table, pa.array(getattr(self, table)))
 
@@ -89,8 +92,10 @@ def __make_reverse_index__(self, verbose: bool) -> None:
             for ind, member in enumerate(getattr(self, table)):
                 index[member["token"]] = ind
 
-            self._token2ind[table] = pd.Series(index.values(), index=index.keys())
+            df = pd.DataFrame({"values": index.values(), "keys": index.keys()})
+            df = df.set_index("keys")["values"]
 
+            self._token2ind[table] = df
         # Decorate (adds short-cut) sample_data with sensor information.
         for record in self.sample_data:
             cs_record = self.get("calibrated_sensor", record["calibrated_sensor_token"])
@@ -119,6 +124,10 @@ def getind(self, table_name: str, token: str) -> int:
         return self._token2ind[table_name][token]
 
 
+# NuScenes = CustomNuScenes
+NuScenes = UnpatchedNuScenes
+
+
 class SampleData(TypedDict):
     timestamp: int
     filename: str
@@ -512,12 +521,13 @@ def _cam2scenes(
                 kls = LidarDataset
             else:
                 raise RuntimeError(f"unsupported sensor_modality {sensor_modality}")
-            samples = pa.array(samples)
+            # samples = pa.array(samples)
             ds_scenes.append(
                 kls(self.data_dir, self.nusc, samples, num_frames=self.num_frames)
             )
 
-        scene_samples = pa.array([sample for scene in scenes for sample in scene])
+        scene_samples = [sample for scene in scenes for sample in scene]
+        # scene_samples = pa.array(scene_samples)
         ds = ConcatDataset(ds_scenes)
         return ds, scene_samples