fix mask

openxrlab · Apr 19, 2024 · e170267 · e170267
1 parent a3372f1
commit e170267
Showing 1 changed file with 10 additions and 11 deletions.
diff --git a/slam/models/conv_onet.py b/slam/models/conv_onet.py
@@ -95,17 +95,16 @@ def pre_precessing(self, cur_frame):
             gt_depth_np = cur_frame.depth
             c2w = cur_frame.get_pose()
             for key, grid in self.grid_c.items():
-                if key != 'grid_coarse':
-                    mask = get_mask_from_c2w(camera=self.camera,
-                                             bound=self.bounding_box,
-                                             c2w=c2w,
-                                             key=key,
-                                             val_shape=grid.val.shape[2:],
-                                             depth_np=gt_depth_np)
-                    mask = torch.from_numpy(mask).permute(
-                        2, 1, 0).unsqueeze(0).unsqueeze(0).repeat(
-                            1, grid.val.shape[1], 1, 1, 1)
-                    self.grid_opti_mask[key] = mask
+                mask = get_mask_from_c2w(camera=self.camera,
+                                         bound=self.bounding_box,
+                                         c2w=c2w,
+                                         key=key,
+                                         val_shape=grid.val.shape[2:],
+                                         depth_np=gt_depth_np)
+                mask = torch.from_numpy(mask).permute(
+                    2, 1, 0).unsqueeze(0).unsqueeze(0).repeat(
+                        1, grid.val.shape[1], 1, 1, 1)
+                self.grid_opti_mask[key] = mask
 
     def get_outputs(self, input) -> Dict[str, Union[torch.Tensor, List]]:
         rays_o = input['rays_o']