meshes

tsr/system.py

@@ -173,31 +173,28 @@ class TSR(BaseModule):
         meshes = []
         for scene_code in scene_codes:
             with torch.no_grad():
-                density = self.renderer.query_triplane(
-                    self.decoder,
-                    scale_tensor(
-                        self.isosurface_helper.grid_vertices.to(scene_codes.device),
-                        self.isosurface_helper.points_range,
-                        (-self.renderer.cfg.radius, self.renderer.cfg.radius),
-                    ),
+                density = self.renderer.query_triplane_volume_density(
+                    self.decoder.to(scene_codes.device),
                     scene_code,
-                )["density_act"]
-            v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold))
-            v_pos = scale_tensor(
-                v_pos,
-                self.isosurface_helper.points_range,
-                (-self.renderer.cfg.radius, self.renderer.cfg.radius),
-            )
-            with torch.no_grad():
+                    resolution
+                ) - threshold
+                v_pos, t_pos_idx = self.isosurface_helper(density)
+                density = None
+                v_pos = v_pos.to(scene_codes.device)
                 color = self.renderer.query_triplane(
-                    self.decoder,
+                    self.decoder.to(scene_codes.device),
                     v_pos,
                     scene_code,
                 )["color"]
-            mesh = trimesh.Trimesh(
-                vertices=v_pos.cpu().numpy(),
-                faces=t_pos_idx.cpu().numpy(),
-                vertex_colors=color.cpu().numpy(),
-            )
-            meshes.append(mesh)
+                v_pos = scale_tensor(
+                    v_pos,
+                    self.isosurface_helper.points_range,
+                    (-self.renderer.cfg.radius, self.renderer.cfg.radius),
+                )
+                mesh = trimesh.Trimesh(
+                    vertices=v_pos.cpu().numpy(),
+                    faces=t_pos_idx.cpu().numpy(),
+                    vertex_colors=color.cpu().numpy(),
+                )
+                meshes.append(mesh)
         return meshes