update mesh occ example

readme.md

@@ -32,6 +32,8 @@ We have included it as a submodule in this repository, but you can also install
 
 ### Usage
 
+**Basics:**
+
 ```python
 import numpy as np
 import trimesh
@@ -60,6 +62,36 @@ distances, face_id, uvw = BVH.signed_distance(points, return_uvw=True, mode='wat
 distances, face_id, uvw = BVH.signed_distance(points, return_uvw=True, mode='raystab') # [N], [N], [N, 3]
 ```
 
+**Robust Mesh Occupancy:**
+
+UDF + flood-fill for possibly non-watertight/single-layer meshes:
+
+```python
+resolution = 512
+device = torch.device('cuda')
+
+BVH = cubvh.cuBVH(vertices, faces)
+
+grid_points = torch.stack(
+    torch.meshgrid(
+        torch.linspace(-1, 1, resolution, device=device),
+        torch.linspace(-1, 1, resolution, device=device),
+        torch.linspace(-1, 1, resolution, device=device),
+        indexing="ij",
+    ), dim=-1,
+) # [N, N, N, 3]
+
+udf, _, _ = BVH.unsigned_distance(grid_points.view(-1, 3), return_uvw=False)
+udf = udf.cpu().numpy().reshape(resolution, resolution, resolution)
+occ = udf < 2 / resolution # tolerance 2 voxels
+
+empty_mask = morphology.flood(occ, (0, 0, 0), connectivity=1) # flood from the corner, which is for sure empty
+occ = ~empty_mask
+```
+Check [`test/robust_occupancy.py`](test/robust_occupancy.py) for more details.
+
+
+**Renderer:**
 
 Example for a mesh normal renderer by `ray_trace`:
 

test/extract_mesh_occupancy.py

+import os
+import zlib
+import time
+import numpy as np
+import torch
+
+import cubvh
+import mcubes
+import trimesh
+import argparse
+
+import kiui
+from kiui.mesh import Mesh
+
+from skimage import morphology
+
+parser = argparse.ArgumentParser()
+parser.add_argument('mesh', type=str)
+parser.add_argument('--res', type=int, default=512)
+opt = parser.parse_args()
+
+device = torch.device('cuda')
+
+points = torch.stack(
+    torch.meshgrid(
+        torch.linspace(-1, 1, opt.res, device=device),
+        torch.linspace(-1, 1, opt.res, device=device),
+        torch.linspace(-1, 1, opt.res, device=device),
+        indexing="ij",
+    ), dim=-1,
+) # [N, N, N, 3]
+
+def run(path):
+    
+    mesh = Mesh.load(path, wotex=True, bound=0.95, device=device)
+
+    t0 = time.time()
+    BVH = cubvh.cuBVH(mesh.v, mesh.f)
+    print('BVH build time:', time.time() - t0)
+
+    # naive sdf
+    # sdf, _, _ = BVH.signed_distance(points.view(-1, 3), return_uvw=False, mode='raystab') # some mesh may not be watertight...
+    # sdf = sdf.cpu().numpy()
+    # occ = (sdf < 0)
+
+    # udf floodfill
+    t0 = time.time()
+    udf, _, _ = BVH.unsigned_distance(points.view(-1, 3), return_uvw=False)
+    print('UDF time:', time.time() - t0)
+    udf = udf.cpu().numpy().reshape(opt.res, opt.res, opt.res)
+    occ = udf < 2 / opt.res # tolerance 2 voxels
+
+    t0 = time.time()
+    empty_mask = morphology.flood(occ, (0, 0, 0), connectivity=1) # flood from the corner, which is for sure empty
+    print('Floodfill time:', time.time() - t0)
+    occ = ~empty_mask
+
+    # # packbits and compress
+    # occ = occ.astype(np.uint8).reshape(-1)
+    # occ = np.packbits(occ)
+    # occ = zlib.compress(occ.tobytes())
+
+    # # save to disk
+    # with open(os.path.basename(path).split('.')[0] + '.bin', 'wb') as f:
+    #     f.write(occ)
+
+    # # uncompress and unpack
+    # occ = zlib.decompress(occ)
+    # occ = np.frombuffer(occ, dtype=np.uint8)
+    # occ = np.unpackbits(occ, count=opt.res**3).reshape(opt.res, opt.res, opt.res)
+
+    # marching cubes
+    occ = occ.reshape(opt.res, opt.res, opt.res)
+    t0 = time.time()
+    verts, faces = mcubes.marching_cubes(occ, 0.5)
+    # smoothed_occ = mcubes.smooth(occ, method='constrained')  # very slow
+    # verts, faces = mcubes.marching_cubes(smoothed_occ, 0)
+    print('MC time:', time.time() - t0)
+
+    verts = verts / (opt.res - 1) * 2 - 1
+    mesh = trimesh.Trimesh(vertices=verts, faces=faces)
+    mesh.export(os.path.basename(path).split('.')[0] + '.ply')
+
+run(opt.mesh)
\ No newline at end of file