Made composed xforms be a sequence. Changed voxel status eval to use inject

Signed-off-by: Christopher Horvath <[email protected]>

Author: blackencino

surface_reconstruction/nksr/nksr_fvdb/coord_xform.py

@@ -80,7 +80,7 @@ def apply_jagged(self, coords: fvdb.JaggedTensor) -> fvdb.JaggedTensor:
         """Transform coordinates from source frame to target frame.
 
         Args:
-            coords: JaggedTensor of shape [N, 3] containing 3D coordinates.
+            coords: JaggedTensor of shape [B, Njagged, 3] containing 3D coordinates.
 
         Returns:
             Transformed coordinates as a JaggedTensor with the same structure.
@@ -101,13 +101,13 @@ def apply(self, coords: torch.Tensor | fvdb.JaggedTensor) -> torch.Tensor | fvdb
         Args:
             coords: 3D coordinates to transform.
                 - torch.Tensor: shape [N, 3] where N is the number of points.
-                - fvdb.JaggedTensor: shape [B][N, 3] where B is batch size and N
+                - fvdb.JaggedTensor: shape [B, Njagged, 3] where B is batch size and Njagged
                   varies per batch element.
 
         Returns:
             Transformed coordinates with the same type and shape as input.
                 - torch.Tensor input -> torch.Tensor output, shape [N, 3]
-                - fvdb.JaggedTensor input -> fvdb.JaggedTensor output, shape [B][N, 3]
+                - fvdb.JaggedTensor input -> fvdb.JaggedTensor output, shape [B, Njagged, 3]
         """
         if isinstance(coords, torch.Tensor):
             return self.apply_tensor(coords)
@@ -288,7 +288,7 @@ def compose(self, other: "CoordXform") -> "CoordXform":
         Returns:
             A new CoordXform representing the composition.
         """
-        return ComposedXform(first=other, second=self)
+        return ComposedXform([other, self])
 
     @overload
     def __matmul__(self, other: "CoordXform") -> "CoordXform": ...
@@ -346,40 +346,54 @@ def __matmul__(
 
 @dataclass(frozen=True)
 class ComposedXform(CoordXform):
-    """Composition of two transformations applied in sequence.
+    """Composition of multiple transformations applied in sequence.
+    The 0th transformation is applied first, followed by the 1st, etc.
 
-    Represents the transformation: coords_out = second(first(coords_in))
     """
 
-    first: CoordXform
-    second: CoordXform
+    xforms: list[CoordXform]
 
     def apply_tensor(self, coords: torch.Tensor) -> torch.Tensor:
-        """Apply first, then second transformation."""
-        return self.second.apply_tensor(self.first.apply_tensor(coords))
+        """Apply all transformations in to tensor coords insequence."""
+        for xform in self.xforms:
+            coords = xform.apply_tensor(coords)
+        return coords
 
     def apply_jagged(self, coords: fvdb.JaggedTensor) -> fvdb.JaggedTensor:
-        """Apply first, then second transformation."""
-        return self.second.apply_jagged(self.first.apply_jagged(coords))
+        """Apply all transformations to jagged tensor coords in sequence."""
+        for xform in self.xforms:
+            coords = xform.apply_jagged(coords)
+        return coords
 
     def apply_bounds_tensor(self, bounds: torch.Tensor) -> torch.Tensor:
-        """Apply first, then second bounds transformation."""
-        return self.second.apply_bounds_tensor(self.first.apply_bounds_tensor(bounds))
+        """Apply all transformations to bounds tensor in sequence."""
+        for xform in self.xforms:
+            bounds = xform.apply_bounds_tensor(bounds)
+        return bounds
 
     def apply_bounds_jagged(self, bounds: fvdb.JaggedTensor) -> fvdb.JaggedTensor:
-        """Apply first, then second bounds transformation."""
-        return self.second.apply_bounds_jagged(self.first.apply_bounds_jagged(bounds))
+        """Apply all transformations to bounds jagged tensor in sequence."""
+        for xform in self.xforms:
+            bounds = xform.apply_bounds_jagged(bounds)
+        return bounds
 
     def inverse(self) -> CoordXform:
-        """Return inverse: (second o first)^-1 = first^-1 o second^-1."""
+        """Return inverse: (all xforms)^-1 = inverse(all xforms) in reverse order."""
         if not self.invertible:
             raise NotImplementedError("Cannot invert: one or both transforms not invertible.")
-        return ComposedXform(self.second.inverse(), self.first.inverse())
+        return ComposedXform([xform.inverse() for xform in reversed(self.xforms)])
+
+    def compose(self, other: "CoordXform") -> "CoordXform":
+        """Compose this transformation with another. Other is applied first, then self."""
+        if isinstance(other, ComposedXform):
+            return ComposedXform(other.xforms + self.xforms)
+        else:
+            return ComposedXform([other] + self.xforms)
 
     @property
     def invertible(self) -> bool:
         """True if both first and second are invertible."""
-        return self.first.invertible and self.second.invertible
+        return all(xform.invertible for xform in self.xforms)
 
 
 @dataclass(frozen=True)
@@ -476,7 +490,8 @@ def compose(self, other: "CoordXform") -> "CoordXform":
                 scale=other.scale * self.scale,
                 translation=other.translation * self.scale + self.translation,
             )
-        return ComposedXform(first=other, second=self)
+        else:
+            return super().compose(other)
 
     @property
     def invertible(self) -> bool:

surface_reconstruction/nksr/nksr_fvdb/sparse_feature_hierarchy.py

@@ -25,7 +25,7 @@ class VoxelStatus(Enum):
     VS_EXIST_CONTINUE = 2
 
 
-def evaluate_voxel_status(target_grid: GridBatch, coarse_grid: GridBatch, fine_grid: GridBatch | None) -> torch.Tensor:
+def evaluate_voxel_status(target_grid: GridBatch, coarse_grid: GridBatch, fine_grid: GridBatch | None) -> JaggedTensor:
     """Compute per-voxel status labels for structure prediction training.
 
     Classifies each voxel in target_grid by comparing against coarse/fine hierarchy levels:
@@ -49,26 +49,32 @@ def evaluate_voxel_status(target_grid: GridBatch, coarse_grid: GridBatch, fine_g
     if coarse_grid.device != device:
         raise ValueError(f"Device not match {device} vs {coarse_grid.device}.")
 
-    if fine_grid is not None and fine_grid.device != device:
-        raise ValueError(f"Device not match {device} vs {fine_grid.device}.")
+    # create a coarse feature jagged tensor, by broadcasting. The coarse grid is filled with
+    # VoxelStatus.VS_EXIST_STOP.
+    coarse_feature = coarse_grid.jagged_like(
+        torch.tensor(VoxelStatus.VS_EXIST_STOP.value, dtype=torch.uint8, device=device).expand(coarse_grid.total_voxels)
+    )
 
-    # Start with a flat tensor of all voxels in the target grid, initialized to NON_EXIST.
-    status = torch.full((target_grid.total_voxels,), VoxelStatus.VS_NON_EXIST.value, dtype=torch.uint8, device=device)
-
-    # Set the intersection of the coarse grid and the target grid to EXIST_STOP.
-    coarse_exist_idx = target_grid.ijk_to_index(coarse_grid.ijk, cumulative=True)
-    coarse_exist_mask = coarse_exist_idx.jdata != -1
-    status[coarse_exist_mask] = VoxelStatus.VS_EXIST_STOP.value
+    # Inject the coarse feature into the target
+    target_feature = coarse_grid.inject_to(
+        dst_grid=target_grid, src=coarse_feature, dst=None, default_value=VoxelStatus.VS_NON_EXIST.value
+    )
 
     # If there is a fine grid, set the intersection of the fine grid and the target grid to EXIST_CONTINUE.
-    # This can overwrite the coarse grid's EXIST_STOP status.
+    # This can overwrite the coarse grid's EXIST_STOP status. It modifies the target_feature in place,
+    # and doesn't overwrite except where the fine coarsened grid exists.
     if fine_grid is not None:
+        if fine_grid.device != device:
+            raise ValueError(f"Device not match {device} vs {fine_grid.device}.")
         fine_coarsened = fine_grid.coarsened_grid(coarsening_factor=2)
-        fine_exist_idx = target_grid.ijk_to_index(fine_coarsened.ijk, cumulative=True)
-        fine_exist_mask = fine_exist_idx.jdata != -1
-        status[fine_exist_mask] = VoxelStatus.VS_EXIST_CONTINUE.value
+        fine_coarsened_feature = fine_coarsened.jagged_like(
+            torch.tensor(VoxelStatus.VS_EXIST_CONTINUE.value, dtype=torch.uint8, device=device).expand(
+                fine_coarsened.total_voxels
+            )
+        )
+        fine_coarsened.inject_to(dst_grid=target_grid, src=fine_coarsened_feature, dst=target_feature)
 
-    return status
+    return target_feature
 
 
 @dataclass(frozen=True)
@@ -270,7 +276,7 @@ def from_point_splatting(
 
         return cls(levels=levels)
 
-    def evaluate_voxel_status(self, target_grid: GridBatch, coarse_depth: int | None = None) -> torch.Tensor:
+    def evaluate_voxel_status(self, target_grid: GridBatch, coarse_depth: int | None = None) -> JaggedTensor:
         """Evaluate voxel status at a given hierarchy depth.
 
         Convenience wrapper around the module-level evaluate_voxel_status function.