Spaces:
Running
Running
| import os | |
| import numpy as np | |
| import taichi as ti | |
| ti.init() | |
| real = ti.f32 | |
| def get_sdf_spatial_grad(): | |
| sdf_sv_fn = "/data/xueyi/diffsim/NeuS/init_box_mesh.npy" | |
| if not os.path.exists(sdf_sv_fn): | |
| sdf_sv_fn = "/home/xueyi/diffsim/NeuS/init_box_mesh.npy" | |
| sdf_sv = np.load(sdf_sv_fn, allow_pickle=True) | |
| scene_sdf = sdf_sv | |
| print(sdf_sv.shape) | |
| res = sdf_sv.shape[0] | |
| sdf_grad = np.zeros((res, res, res, 3), dtype=np.float32) | |
| # scene_sdf = ti.field(dtype=real, shape=(self.sdf_res, self.sdf_res, self.sdf_res)) | |
| for i_x in range(res): | |
| print(f"Start processing i_x : {i_x}") | |
| for i_y in range(res): | |
| print(f"Start processing i_x : {i_x}, i_y : {i_y}") | |
| for i_z in range(res): | |
| cur_grad = np.zeros((3,), dtype=np.float32) | |
| if i_x > 0 and i_x < res - 1: | |
| cur_grad[0] = (scene_sdf[i_x + 1, i_y, i_z] - scene_sdf[i_x - 1, i_y, i_z]) / 2.0 | |
| elif i_x == 0: | |
| cur_grad[0] = scene_sdf[i_x + 1, i_y, i_z] - scene_sdf[i_x, i_y, i_z] | |
| elif i_x == res - 1: | |
| cur_grad[0] = scene_sdf[i_x, i_y, i_z] - scene_sdf[i_x - 1, i_y, i_z] | |
| if i_y > 0 and i_y < res - 1: | |
| cur_grad[1] = (scene_sdf[i_x, i_y + 1, i_z] - scene_sdf[i_x, i_y - 1, i_z]) / 2.0 | |
| elif i_y == 0: | |
| cur_grad[1] = scene_sdf[i_x, i_y + 1, i_z] - scene_sdf[i_x, i_y, i_z] | |
| elif i_y == res - 1: | |
| cur_grad[1] = scene_sdf[i_x, i_y, i_z] - scene_sdf[i_x, i_y - 1, i_z] | |
| if i_z > 0 and i_z < res - 1: | |
| cur_grad[2] = (scene_sdf[i_x, i_y, i_z + 1] - scene_sdf[i_x, i_y, i_z - 1]) / 2.0 | |
| elif i_z == 0: | |
| cur_grad[2] = scene_sdf[i_x, i_y, i_z + 1] - scene_sdf[i_x, i_y, i_z] | |
| elif i_z == res - 1: | |
| cur_grad[2] = scene_sdf[i_x, i_y, i_z] - scene_sdf[i_x, i_y, i_z - 1] | |
| sdf_grad[i_x, i_y, i_z, :] = cur_grad[:] | |
| sdf_grad_sv_fn = "/home/xueyi/diffsim/NeuS/init_box_mesh_sdf_grad.npy" | |
| np.save(sdf_grad_sv_fn, sdf_grad) | |
| if __name__ == '__main__': | |
| get_sdf_spatial_grad() # |