Upload 40 files

requirements.txt

@@ -1,7 +1,7 @@
 omegaconf==2.3.0
 Pillow==10.1.0
 einops==0.7.0
-scikit-image
+git+https://github.com/cocktailpeanut/torchmcubes.git
 transformers==4.35.0
 trimesh==4.0.5
 onnxruntime

tsr/models/isosurface.py

@@ -3,7 +3,7 @@ from typing import Callable, Optional, Tuple
 import numpy as np
 import torch
 import torch.nn as nn
-from skimage import measure
+from torchmcubes import marching_cubes
 
 
 class IsosurfaceHelper(nn.Module):
@@ -18,6 +18,7 @@ class MarchingCubeHelper(IsosurfaceHelper):
     def __init__(self, resolution: int) -> None:
         super().__init__()
         self.resolution = resolution
+        self.mc_func: Callable = marching_cubes
         self._grid_vertices: Optional[torch.FloatTensor] = None
 
     @property
@@ -40,10 +41,7 @@ class MarchingCubeHelper(IsosurfaceHelper):
         level: torch.FloatTensor,
     ) -> Tuple[torch.FloatTensor, torch.LongTensor]:
         level = -level.view(self.resolution, self.resolution, self.resolution)
-        level_np = level.detach().cpu().numpy()
-        v_pos_np, t_pos_idx_np, _, _ = measure.marching_cubes(level_np, level=0.0)
-        v_pos = torch.from_numpy(v_pos_np.astype(np.float32)).to(level.device)
-        t_pos_idx = torch.from_numpy(t_pos_idx_np.astype(np.int64)).to(level.device)
+        v_pos, t_pos_idx = self.mc_func(level.detach(), 0.0)
         v_pos = v_pos[..., [2, 1, 0]]
         v_pos = v_pos / (self.resolution - 1.0)
         return v_pos, t_pos_idx