voxel: add grid floor

d4l3k · Oct 19, 2023 · d66858b · d66858b
1 parent b8fd677
commit d66858b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/autolabel.py b/autolabel.py
@@ -126,7 +126,7 @@ def forward(self, img: torch.Tensor) -> torch.Tensor:
 
 TASKS = [
     LabelDet(),
-    # LabelSemSeg(),
+    LabelSemSeg(),
 ]
 
 
@@ -185,7 +185,8 @@ def get_task_path(task: nn.Module) -> str:
         cam_data = {}
         for cam, frames in batch.color.items():
             squashed = frames[idxs].squeeze(1)
-            cam_data[cam] = task(squashed)
+            with torch.no_grad():
+                cam_data[cam] = task(squashed)
 
         for j, i in enumerate(idxs):
             frame_data = {}

diff --git a/torchdrive/tasks/voxel.py b/torchdrive/tasks/voxel.py
@@ -245,6 +245,13 @@ def forward(
         )[:, 1:]
 
         grid = grid.permute(0, 1, 4, 3, 2)
+
+        # set floor
+        # this helps avoid degenerate cases due to dynamic objects, blur and
+        # reflections
+        grid = grid.clone()
+        grid[:, :, :, :, 0] = 1
+
         feat_grid = feat_grid.permute(0, 1, 4, 3, 2)
         # agrid, _ = axis_grid(grid)
         # grid = (grid+agrid).clamp(min=0, max=1)