#positions = scale_tensor( # positions, (-self.cfg.radius, self.cfg.radius), (-1, 1) #)

tsr/models/nerf_renderer.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, field
-from typing import Dict
+from typing import Dict, Optional
 
 import torch
 import torch.nn.functional as F
@@ -37,7 +37,59 @@ class TriplaneNeRFRenderer(BaseModule):
             chunk_size >= 0
         ), "chunk_size must be a non-negative integer (0 for no chunking)."
         self.chunk_size = chunk_size
+     def make_step_grid(self,device, resolution: int, chunk_size: int = 32):
+        coords = torch.linspace(-1.0, 1.0, resolution, device = device)
+        x, y, z = torch.meshgrid(coords[0:chunk_size], coords, coords, indexing="ij")
+        x = x.reshape(-1, 1)
+        y = y.reshape(-1, 1)
+        z = z.reshape(-1, 1)
+        verts = torch.cat([x, y, z], dim = -1).view(-1, 3)
+        indices2D: torch.Tensor = torch.stack(
+            (verts[..., [0, 1]], verts[..., [0, 2]], verts[..., [1, 2]]),
+            dim=-3,
+        )
+        return indices2D
+
+    def query_triplane_volume_density(self, decoder: torch.nn.Module, triplane: torch.Tensor, resolution: int, sample_count: int = 1024 * 1024 * 4) -> torch.Tensor:
+        layer_count = sample_count // (resolution * resolution)
+        out_list = self.do_query_triplane_volume_density(decoder, triplane, resolution, layer_count)
+        return get_activation(self.cfg.density_activation)(
+            out_list.view([resolution * resolution * resolution, 1]) + self.cfg.density_bias
+        )
+    def do_query_triplane_volume_density(self, decoder: torch.nn.Module, triplane: torch.Tensor, resolution: int, layer_count: int) -> torch.Tensor:
+        step = 2.0 * layer_count / (resolution - 1)
+        indices2D = self.make_step_grid(triplane.device, resolution, layer_count)
+
+        out_list = torch.zeros([resolution, resolution * resolution, 1], device = triplane.device
+                               )
+        for i in range(0, resolution, layer_count):
+            if i + layer_count > resolution:
+                layer_count = resolution - i
+                indices2D = indices2D[..., :resolution * resolution * layer_count, :]
+            density_step = self.sample_step_triplane_volume_density(decoder, triplane, indices2D)
+            # todo directly march cube
+            out_list[i:i + layer_count] = density_step.view([layer_count, resolution * resolution, 1])
+            #out_list.append(net_out['density'])
+            indices2D.transpose(1, 2)[0, 0] += step
+            indices2D.transpose(1, 2)[1, 0] += step
+
+        return out_list
+    def sample_step_triplane_volume_density(self, decoder, triplane, indices2D):
+        out: torch.Tensor = F.grid_sample(
+            rearrange(triplane, "Np Cp Hp Wp -> Np Cp Hp Wp", Np=3),
+            rearrange(indices2D, "Np N Nd -> Np () N Nd", Np=3),
+            align_corners=False,
+            mode="bilinear",
+        )
+        if self.cfg.feature_reduction == "concat":
+            out = rearrange(out, "Np Cp () N -> N (Np Cp)", Np=3)
+        elif self.cfg.feature_reduction == "mean":
+            out = reduce(out, "Np Cp () N -> N Cp", Np=3, reduction="mean")
+        else:
+            raise NotImplementedError
 
+        net_out: Dict[str, torch.Tensor] = decoder(out)
+        return net_out['density']
     def query_triplane(
         self,
         decoder: torch.nn.Module,
@@ -49,9 +101,9 @@ class TriplaneNeRFRenderer(BaseModule):
 
         # positions in (-radius, radius)
         # normalized to (-1, 1) for grid sample
-        positions = scale_tensor(
-            positions, (-self.cfg.radius, self.cfg.radius), (-1, 1)
-        )
+        #positions = scale_tensor(
+        #    positions, (-self.cfg.radius, self.cfg.radius), (-1, 1)
+        #)
 
         def _query_chunk(x):
             indices2D: torch.Tensor = torch.stack(