use 32 feats for BEV/voxel by default. Add nfeats arg to PCA block (#27)

Author: striest

config/image_processing/radiov3_siglip2_pca.yaml

@@ -8,4 +8,5 @@
 - 
     type: pca
     args:
-        fp: physics_atv_visual_mapping/pca/radio_v3b_siglip2.pt
+        fp: physics_atv_visual_mapping/pca/radio_v3b_siglip2.pt
+        n_features: 32

config/terrain_estimation/voxel_to_bev.yaml

@@ -88,7 +88,7 @@
     type: terrain_aware_bev_feature_splat
     args:
         metainfo_key: vfm
-        n_features: 16
+        n_features: 32
         terrain_layer: terrain
         terrain_mask_layer: min_elevation_filtered_inflated_mask
         overhang: 2.0

physics_atv_visual_mapping/image_processing/processing_blocks/pca.py

@@ -11,7 +11,7 @@ class PCABlock(ImageProcessingBlock):
     Block that applies a precomputed PCA to the image
     """
 
-    def __init__(self, fp, models_dir, device):
+    def __init__(self, fp, models_dir, n_features=-1, device='cuda'):
         full_fp = os.path.join(models_dir, fp)
         pca = torch.load(full_fp, weights_only=False)
 
@@ -28,6 +28,10 @@ def __init__(self, fp, models_dir, device):
             "V": pca["V"].to(device),
         }
 
+        self.n_features = n_features
+
+        assert self.n_features <= self.pca["V"].shape[-1] 
+
     def run(self, image, intrinsics, image_orig):
         _pmean = self.pca["mean"].view(1, 1, -1)
         _pv = self.pca["V"].unsqueeze(0)
@@ -40,11 +44,15 @@ def run(self, image, intrinsics, image_orig):
             image.shape[0], _pv.shape[-1], image.shape[2], image.shape[3]
         )
 
+        if self.n_features >= 0:
+            img_out = img_out[:, :self.n_features]
+
         return img_out, intrinsics
 
     @property
     def output_feature_keys(self):
+        N = self.pca["V"].shape[-1] if self.n_features == -1 else self.n_features
         return FeatureKeyList(
-            label=[f"{self.base_label}_{i}" for i in range(self.pca["V"].shape[-1])],
-            metainfo=["vfm" for i in range(self.pca["V"].shape[-1])]
+            label=[f"{self.base_label}_{i}" for i in range(N)],
+            metainfo=["vfm" for i in range(N)]
         )

scripts/offline_processing/multiproc.py

@@ -27,20 +27,34 @@
     for rdir in sorted(run_dirs):
         print('\t' + rdir)
 
-    base_cmd = "python3 get_voxel_inpainting_supervision.py --run_dir {}"
+    # base_cmd = "python3 get_voxel_inpainting_supervision.py --run_dir {}"
+
+    base_cmd = "python3 sdf_to_traj.py --run_dir {}"
+    base_cmd2 = "python3 sdf_to_traj.py --run_dir {} --bev_dir bev_map_inpaint_reduce"
 
     success_dirs = []
     fail_dirs = []
 
     for ri, run_dir in enumerate(run_dirs):
         print("Proc {} ({}/{})".format(run_dir, ri+1, len(run_dirs)))
 
-        # cmd = base_cmd.format(args.config_fp, run_dir)
-        cmd = base_cmd.format(run_dir)
+        # cmd = base_cmd.format(run_dir)
+
+        # res = subprocess.run(cmd.split(" "))
+
+        # if res.returncode == 0:
+        #     success_dirs.append(run_dir)
+        # else:
+        #     fail_dirs.append(run_dir)
+
+
+        cmd1 = base_cmd.format(run_dir)
+        cmd2 = base_cmd2.format(run_dir)
 
-        res = subprocess.run(cmd.split(" "))
+        res = subprocess.run(cmd1.split(" "))
+        res2 = subprocess.run(cmd2.split(" "))
 
-        if res.returncode == 0:
+        if (res.returncode == 0) and (res2.returncode == 0):
             success_dirs.append(run_dir)
         else:
             fail_dirs.append(run_dir)

scripts/offline_processing/sdf_to_traj.py

@@ -57,8 +57,13 @@
         mindists = torch.linalg.norm(coords.unsqueeze(1) - poses.unsqueeze(0), dim=-1).min(dim=1)[0]
         mindists = mindists.reshape(*metadata.N)
 
-        bev_grid.bev_grid.feature_keys += FeatureKeyList(label=['dist_to_traj'], metainfo=['spatial'])
-        bev_grid.bev_grid.data = torch.cat([bev_grid.bev_grid.data, mindists.unsqueeze(-1)], dim=-1)
+        ## remove spatial if there
+        idxs = [i for i,k in enumerate(bev_grid.bev_grid.feature_keys.label) if k != 'dist_to_traj']
+        _fks = bev_grid.bev_grid.feature_keys[idxs]
+        _data = bev_grid.bev_grid.data[..., idxs]
+
+        bev_grid.bev_grid.feature_keys = _fks + FeatureKeyList(label=['dist_to_traj'], metainfo=['spatial'])
+        bev_grid.bev_grid.data = torch.cat([_data, mindists.unsqueeze(-1)], dim=-1)
 
         bev_grid.to_kitti(bev_dir, i)