diff --git a/requirements.txt b/requirements.txt index 5fe390d..99f5298 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ omegaconf==2.3.0 Pillow==10.1.0 einops==0.7.0 -git+https://github.com/tatsy/torchmcubes.git +scikit-image transformers==4.35.0 trimesh==4.0.5 rembg diff --git a/tsr/models/isosurface.py b/tsr/models/isosurface.py index 076321b..14a2855 100644 --- a/tsr/models/isosurface.py +++ b/tsr/models/isosurface.py @@ -3,7 +3,6 @@ import numpy as np import torch import torch.nn as nn -from torchmcubes import marching_cubes class IsosurfaceHelper(nn.Module): @@ -42,11 +41,12 @@ def forward( level: torch.FloatTensor, ) -> Tuple[torch.FloatTensor, torch.LongTensor]: level = -level.view(self.resolution, self.resolution, self.resolution) - try: - v_pos, t_pos_idx = self.mc_func(level.detach(), 0.0) - except AttributeError: - print("torchmcubes was not compiled with CUDA support, use CPU version instead.") - v_pos, t_pos_idx = self.mc_func(level.detach().cpu(), 0.0) - v_pos = v_pos[..., [2, 1, 0]] + v_pos, t_pos_idx, _, __ = measure.marching_cubes( + (level.detach().cpu() if level.is_cuda else level.detach()).numpy(), + 0.0) # self.mc_func(level.detach(), 0.0) + v_pos = torch.from_numpy(v_pos.copy()).type(torch.FloatTensor).to(level.device) + t_pos_idx = torch.from_numpy(t_pos_idx.copy()).type(torch.LongTensor).to(level.device) + v_pos = v_pos[..., [0, 1, 2]] + t_pos_idx = t_pos_idx[..., [1, 0, 2]] v_pos = v_pos / (self.resolution - 1.0) - return v_pos.to(level.device), t_pos_idx.to(level.device) + return v_pos, t_pos_idx