nuscenes: lidar loading fixes + metrics fixes

d4l3k · Sep 15, 2023 · a2dcff8 · a2dcff8
1 parent 6f0b2c3
commit a2dcff8
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/torchdrive/datasets/nuscenes_dataset.py b/torchdrive/datasets/nuscenes_dataset.py
@@ -279,7 +279,7 @@ def __getitem__(self, idx: int) -> object:
         pcl = LidarPointCloud.from_file(
             os.path.join(self.dataroot, sample_data["filename"])
         )
-        return pcl.points
+        return torch.from_numpy(pcl.points)
 
 
 class NuscenesDataset(Dataset):
@@ -292,7 +292,9 @@ class NuscenesDataset(Dataset):
         CamTypes.CAM_BACK_RIGHT: [CamTypes.CAM_BACK, CamTypes.CAM_FRONT_RIGHT],
     }
 
-    def __init__(self, data_dir: str, version: str = "v1.0-trainval") -> None:
+    def __init__(
+        self, data_dir: str, version: str = "v1.0-trainval", lidar: bool = False
+    ) -> None:
         self.data_dir = data_dir
         self.version = version
         self.nusc = NuScenes(version=version, dataroot=data_dir, verbose=True)
@@ -304,7 +306,10 @@ def __init__(self, data_dir: str, version: str = "v1.0-trainval") -> None:
             CamTypes.CAM_BACK_LEFT,
             CamTypes.CAM_BACK_RIGHT,
         ]
-        self.sensor_types = self.cam_types + [SensorTypes.LIDAR_TOP]
+        self.sensor_types = self.cam_types
+        if lidar:
+            self.sensor_types = self.sensor_types + [SensorTypes.LIDAR_TOP]
+        self.lidar = lidar
         self.cameras: List[str] = list(self.CAMERA_OVERLAP.keys())
 
         # Organize all the sample_data into scenes by camera type
@@ -399,7 +404,7 @@ def _getitem(self, idx: int) -> Optional[Batch]:
             colors[cam] = torch.stack(cam_colors, dim=0)
             masks[cam] = sample_dict["mask"]
 
-        lidar = data[SensorTypes.LIDAR_TOP]
+        lidar = data[SensorTypes.LIDAR_TOP] if self.lidar else None
 
         return Batch(
             weight=weight.float(),

diff --git a/torchdrive/tasks/voxel.py b/torchdrive/tasks/voxel.py
@@ -596,10 +596,12 @@ def _losses(
                         add_text=False,
                     )
                     ctx.add_figure("semantic/confusion_matrix", fig)
+                    self.semantic_confusion_matrix.reset()
                 if ctx.log_text:
                     ctx.add_scalar(
                         "semantic/accuracy", self.semantic_accuracy.compute()
                     )
+                    self.semantic_accuracy.reset()
 
             del voxel_depth
             del semantic_vel