OcMesher release

.gitignore

@@ -0,0 +1,3 @@
+__pycache__
+lib
+results

LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2023, Princeton Vision & Learning Lab
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -0,0 +1,53 @@
+## OcMesher: View-Dependent Octree-based Mesh Extraction in Unbounded Scenes for Procedural Synthetic Data
+
+Implementation source-code for <it>OcMesher</it>, which extracts a mesh for an unbounded scene represented by signed distance functions (SDFs). Even though the scene is unbounded, the mesh is memory-efficient, and highly detailed from a given set of camera views. OcMesher is used by default in [Infinigen](https://github.com/princeton-vl/infinigen) to speed up video generation, improve rendering quality, and export terrain meshes to external simulators.
+
+<img src=".github/OcMesher.png" width='1000'>
+
+If you use OcMesher in your work, please cite our academic paper:
+
+<h3 align="center">
+    <a href=".">
+        View-Dependent Octree-based Mesh Extraction in Unbounded Scenes for Procedural Synthetic Data
+    </a>
+</h3>
+<p align="center">
+    <a href="https://mazeyu.github.io/">Zeyu Ma</a>, 
+    <a href="http://araistrick.com/">Alexander Raistrick</a>, 
+    <a href="https://www.lahavlipson.com/">Lahav Lipson</a>, 
+    <a href="http://www.cs.princeton.edu/~jiadeng">Jia Deng</a><br>
+</p>
+
+```
+@article{ocmesher2023view,
+  title={View-Dependent Octree-based Mesh Extraction in Unbounded Scenes for Procedural Synthetic Data},
+  author={Ma, Zeyu and Raistrick, Alexander and Lipson, Lahav and Deng, Jia},
+  year={2023}
+}
+```
+
+Please view the video [here](https://youtu.be/YA1c5L0Ncuw) for more qualitative results
+
+## Getting Started
+
+:bulb: Note: OcMesher is installed by default in Infinigen as of v1.2.0 - if you wish to use OcMesher with Infinigen please follow the Installation instructions on the infinigen repo. Use the instructions below only if you want a standalone installation & demo. 
+
+### Standalone Installation
+
+```
+git clone https://github.com/princeton-vl/OcMesher.git
+cd OcMesher
+bash install.sh
+conda create --name ocmesher python=3.10
+conda activate ocmesher
+pip install -r requirements.txt
+```
+
+
+### Demo
+
+```
+python demo.py
+```
+
+This example uses one camera and the Perlin Noise from the Python library `vnoise` and outputs the resulting mesh in `results/demo.obj`.

demo.py

@@ -0,0 +1,38 @@
+# Copyright (c) Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Zeyu Ma
+
+import os
+import sys
+import numpy as np
+from ocmesher import OcMesher
+
+import vnoise
+noise = vnoise.Noise()
+
+def f(XYZ):
+    scale = 2
+    h = noise.noise2(XYZ[:, 0] / scale, XYZ[:, 1] / scale, grid_mode=False, octaves=4)
+
+    return XYZ[:, 2] - h
+
+
+cam_poses = [np.array([
+    [1, 0, 0, 0],
+    [0, 0, 1, 0],
+    [0, -1, 0, 3],
+    [0, 0, 0, 1],
+])]
+Ks = [np.array([
+    [2000, 0, 640],
+    [0, 2000, 360],
+    [0, 0, 1]
+])]
+Hs = [720]
+Ws = [1280]
+
+mesher = OcMesher((cam_poses, Ks, Hs, Ws), pixels_per_cube=16)
+meshes, in_view_tags = mesher([f])
+os.makedirs("results", exist_ok=True)
+meshes[0].export("results/demo.obj")

install.sh

@@ -0,0 +1,27 @@
+#!/bin/bash
+
+shopt -s expand_aliases
+set -e
+
+OS=$(uname -s)
+ARCH=$(uname -m)
+
+if [ "${OS}" = "Linux" ]; then
+    alias gx1="g++ -O3 -c -fpic -fopenmp "
+    alias gx2="g++ -O3 -shared -fopenmp "
+elif [ "${OS}" = "Darwin" ]; then
+    if [ "${ARCH}" = "arm64" ]; then
+        compiler="/opt/homebrew/opt/llvm/bin/clang++"
+    else
+        compiler="/usr/local/opt/llvm/bin/clang++"
+    fi
+    alias gx1="${compiler} -O3 -c -fpic -fopenmp "
+    alias gx2="${compiler} -O3 -shared -fopenmp "
+else
+    echo "Unsupported OS"
+    exit -1
+fi
+
+mkdir -p ocmesher/lib
+gx1 -o ocmesher/lib/core.o ocmesher/source/core.cpp
+gx2 -o ocmesher/lib/core.so ocmesher/lib/core.o

ocmesher/__init__.py

@@ -0,0 +1 @@
+from .core import OcMesher

ocmesher/core.py

@@ -0,0 +1,230 @@
+# Copyright (c) Princeton University.
+# This source code is licensed under the BSD 3-Clause license found in the LICENSE file in the root directory of this source tree.
+
+# Authors: Zeyu Ma
+
+from pathlib import Path
+import sys
+import gin
+import numpy as np
+import trimesh
+from tqdm import tqdm
+
+from .utils.interface import AC, POINTER, AsDouble, AsFloat, AsInt, AsBool, c_bool, c_double, c_float, c_int32, load_cdll, register_func
+from .utils.timer import Timer
+
+@gin.configurable
+class OcMesher:
+    def __init__(self,
+        cameras,
+        pixels_per_cube=8,
+        inv_scale=10,
+        min_dist=1,
+        center=(0, 0, 0), size=1e3,
+        memory_limit_mb=1000,
+        bisection_iters=15,
+        enclosed=True,
+        simplify_occluded=True,
+        visible_relax_iter=2,
+        coarse_count=500000,
+    ):
+        dll = load_cdll(str(Path(__file__).parent.resolve()/"lib"/"core.so"))
+        self.float_type = c_double
+        self.np_float_type = np.float64
+        self.AF = AsDouble
+        self.sdf_float_type = c_float
+        self.sdf_np_float_type = np.float32
+        self.sdf_AF = AsFloat
+
+        self.memory_limit_mb = memory_limit_mb
+
+        cam_poses, Ks, Hs, Ws = cameras
+        self.n_cameras = len(cam_poses)
+        self.cameras = np.zeros(23 * self.n_cameras, dtype=self.np_float_type)
+        for i in range(self.n_cameras):
+            self.cameras[23 * i: 23 * (i+1)] = np.concatenate([
+                np.linalg.inv(cam_poses[i])[:3, :4].reshape(-1),
+                Ks[i].reshape(-1), [Hs[i]], [Ws[i]]
+            ]).astype(self.np_float_type)
+
+        self.inview_pixels_per_cube = self.np_float_type(pixels_per_cube)
+        self.inv_scale = self.np_float_type(inv_scale)
+        self.min_dist = self.np_float_type(min_dist)
+
+        self.center = np.array(center, self.np_float_type)
+        self.size = self.np_float_type(size)
+
+        self.bisection_iters = bisection_iters
+        self.enclosed = enclosed
+        self.simplify_occluded = simplify_occluded
+        self.visible_relax_iter = visible_relax_iter
+        self.coarse_count = coarse_count
+
+        register_func(self, dll, "run_coarse", [
+            POINTER(self.float_type), self.float_type,
+            c_int32, POINTER(self.float_type),
+            self.float_type, self.float_type, self.float_type,
+            c_int32, c_int32, c_int32,
+        ], c_int32)
+        register_func(self, dll, "fine_group", [], c_int32)
+        register_func(self, dll, "fine_iteration", [POINTER(self.sdf_float_type)], c_int32)
+        register_func(self, dll, "fine_iteration_output", [POINTER(self.float_type)])
+        register_func(self, dll, "vis_filter", [c_bool, c_int32], c_int32)
+        register_func(self, dll, "final_iteration", [], c_int32)
+        register_func(self, dll, "final_iteration_occluded", [], c_int32)
+        register_func(self, dll, "final_iteration2", [POINTER(self.float_type)])
+        register_func(self, dll, "final_iteration3", [POINTER(self.sdf_float_type)], c_int32)
+        register_func(self, dll, "final_iteration3_occluded", [POINTER(self.sdf_float_type)])
+        register_func(self, dll, "final_remaining", [POINTER(c_int32)])
+        register_func(self, dll, "get_verts_center", [c_int32, POINTER(self.float_type)])
+        register_func(self, dll, "update_verts", [c_int32, POINTER(self.sdf_float_type), POINTER(self.sdf_float_type), POINTER(self.float_type)])
+        register_func(self, dll, "get_lr_verts", [c_int32, POINTER(self.float_type), POINTER(self.float_type)])
+        register_func(self, dll, "finalize_verts", [c_int32, POINTER(self.sdf_float_type), POINTER(self.sdf_float_type), POINTER(self.float_type)])
+        register_func(self, dll, "construct_faces", [c_int32, POINTER(self.float_type), POINTER(c_int32)])
+        register_func(self, dll, "get_extra_verts_center", [POINTER(self.float_type), POINTER(self.float_type)])
+        register_func(self, dll, "update_extra_verts", [
+            POINTER(self.sdf_float_type), POINTER(self.sdf_float_type),
+            POINTER(self.sdf_float_type), POINTER(self.sdf_float_type),
+            POINTER(self.float_type), POINTER(self.float_type),
+        ])
+        register_func(self, dll, "get_lr_extra_verts", [POINTER(self.float_type), POINTER(self.float_type), POINTER(self.float_type), POINTER(self.float_type)])
+        register_func(self, dll, "finalize_extra_verts", [
+            POINTER(self.sdf_float_type), POINTER(self.sdf_float_type), POINTER(self.float_type),
+            POINTER(self.sdf_float_type), POINTER(self.sdf_float_type), POINTER(self.float_type),
+        ])
+        register_func(self, dll, "get_faces", [POINTER(c_int32)])
+        register_func(self, dll, "get_in_view_tag", [c_int32, POINTER(c_bool)])
+
+
+    def kernel_caller(self, kernels, XYZ_all):
+        n_XYZ = len(XYZ_all)
+        if n_XYZ == 0: return np.zeros((0, len(kernels)), dtype=self.sdf_np_float_type)
+        step = 10000000
+        sdfs = []
+        for i in range(0, n_XYZ, step):
+            XYZ = XYZ_all[i: i+step]
+            sdfs_i = []
+            if self.enclosed:
+                out_bound = np.zeros(len(XYZ), dtype=bool)
+                for i in range(3):
+                    out_bound |= XYZ[:, i] <= self.center[i] - self.size / 2
+                    out_bound |= XYZ[:, i] >= self.center[i] + self.size / 2
+            for kernel in kernels:
+                sdf = kernel(XYZ)
+                if self.enclosed: sdf[out_bound] = 1
+                sdfs_i.append(sdf)
+            sdfs.append(np.stack(sdfs_i, -1).astype(self.sdf_np_float_type))
+        return np.concatenate(sdfs, 0)
+
+    def __call__(self, kernels):
+        n_elements = len(kernels)
+        # octree only considering cameras, not sdf
+        with Timer("coarse step part1"):
+            n_blocks = self.run_coarse(
+                self.AF(self.center), self.size,
+                self.n_cameras, self.AF(self.cameras),
+                self.inview_pixels_per_cube, 
+                self.inv_scale, self.min_dist,
+                self.coarse_count, self.memory_limit_mb, n_elements
+            )
+        # start considering sdf
+        with Timer("coarse step part2"), tqdm(total=n_blocks) as pbar:
+            while True:
+                inc = self.fine_group()
+                if inc == 0: break
+                pbar.update(inc)
+                n = self.fine_iteration(POINTER(self.sdf_float_type)())
+                while n > 0:
+                    positions = AC(np.zeros((n, 3), dtype=self.np_float_type))
+                    self.fine_iteration_output(self.AF(positions))
+                    sdf = AC(self.kernel_caller(kernels, positions).min(axis=-1))
+                    n = self.fine_iteration(self.sdf_AF(sdf))
+        with Timer("filter visible blocks"):
+            n_vis_block = self.vis_filter(self.simplify_occluded, self.visible_relax_iter)
+
+        with Timer("fine step"), tqdm(total=n_vis_block) as pbar:
+            nv = np.zeros(1, dtype=np.int32)
+            while True:
+                n = self.final_iteration(AsInt(nv))
+                if n == 0: break
+                positions = AC(np.zeros((n, 3), dtype=self.np_float_type))
+                self.final_iteration2(self.AF(positions))
+                sdf = AC(self.kernel_caller(kernels, positions))
+                inc = self.final_iteration3(self.sdf_AF(sdf))
+                pbar.update(inc)
+            n = self.final_iteration_occluded(AsInt(nv))
+            if n != 0:
+                positions = AC(np.zeros((n, 3), dtype=self.np_float_type))
+                self.final_iteration2(self.AF(positions))
+                sdf = AC(self.kernel_caller(kernels, positions))
+                self.final_iteration3_occluded(self.sdf_AF(sdf))
+            nv = np.zeros(n_elements, dtype=np.int32)
+            self.final_remaining(AsInt(nv))
+            del positions, sdf
+
+        with Timer("construct mesh"):
+            meshes = []
+            in_view_tags = []
+            for e in range(n_elements):
+                k_e = kernels[e:e+1]
+                centers = np.zeros((nv[e], 3), dtype=self.np_float_type)
+                self.get_verts_center(e, self.AF(centers))
+                center_sdf = self.kernel_caller(k_e, centers)
+                cubes = AC(np.zeros((nv[e] * 8, 3), dtype=self.np_float_type))
+                self.update_verts(e, POINTER(self.sdf_float_type)(), POINTER(self.sdf_float_type)(), self.AF(cubes))
+                for _ in tqdm(range(self.bisection_iters)):
+                    sdf = self.kernel_caller(k_e, cubes)
+                    self.update_verts(e, self.sdf_AF(AC(sdf)), self.sdf_AF(AC(center_sdf)), self.AF(cubes))
+                cubes_r = AC(np.zeros((nv[e] * 8, 3), dtype=self.np_float_type))
+                self.get_lr_verts(e, self.AF(cubes), self.AF(cubes_r))
+                sdf_l = self.kernel_caller(k_e, cubes)
+                sdf_r = self.kernel_caller(k_e, cubes_r)
+                del cubes, cubes_r, centers, center_sdf
+                vertices = np.zeros((nv[e], 3), dtype=self.np_float_type)
+                self.finalize_verts(e, self.sdf_AF(sdf_l), self.sdf_AF(sdf_r), self.AF(vertices))
+                del sdf_l, sdf_r
+                cnts = np.zeros(3, dtype=np.int32)
+                self.construct_faces(e, self.AF(vertices), AsInt(cnts))
+                nve, nvf, nf = cnts
+                edge_vertices_c = AC(np.zeros((nve, 3), dtype=self.np_float_type))
+                face_vertices_c = AC(np.zeros((nvf, 3), dtype=self.np_float_type))
+                self.get_extra_verts_center(self.AF(edge_vertices_c), self.AF(face_vertices_c))
+                ecenter_sdf = self.kernel_caller(k_e, edge_vertices_c)
+                fcenter_sdf = self.kernel_caller(k_e, face_vertices_c)
+                edge_vertices_lr = AC(np.zeros((nve * 2, 3), dtype=self.np_float_type))
+                face_vertices_lr = AC(np.zeros((nvf * 4, 3), dtype=self.np_float_type))
+                self.update_extra_verts(
+                    POINTER(self.sdf_float_type)(), POINTER(self.sdf_float_type)(),
+                    POINTER(self.sdf_float_type)(), POINTER(self.sdf_float_type)(),
+                    self.AF(edge_vertices_lr), self.AF(face_vertices_lr),
+                )
+                for _ in range(self.bisection_iters):
+                    e_sdf = self.kernel_caller(k_e, edge_vertices_lr)
+                    f_sdf = self.kernel_caller(k_e, face_vertices_lr)
+                    self.update_extra_verts(
+                        self.sdf_AF(e_sdf), self.sdf_AF(f_sdf),
+                        self.sdf_AF(ecenter_sdf), self.sdf_AF(fcenter_sdf),
+                        self.AF(edge_vertices_lr), self.AF(face_vertices_lr),
+                    )
+                del edge_vertices_c, face_vertices_c, ecenter_sdf, fcenter_sdf
+                edge_vertices_r = AC(np.zeros((nve * 2, 3), dtype=self.np_float_type))
+                face_vertices_r = AC(np.zeros((nvf * 4, 3), dtype=self.np_float_type))
+                self.get_lr_extra_verts(self.AF(edge_vertices_lr), self.AF(edge_vertices_r), self.AF(face_vertices_lr), self.AF(face_vertices_r))
+                esdf_l = self.kernel_caller(k_e, edge_vertices_lr)
+                esdf_r = self.kernel_caller(k_e, edge_vertices_r)
+                fsdf_l = self.kernel_caller(k_e, face_vertices_lr)
+                fsdf_r = self.kernel_caller(k_e, face_vertices_r)
+                del edge_vertices_lr, edge_vertices_r, face_vertices_lr, face_vertices_r
+                edge_vertices = np.zeros((nve, 3), dtype=self.np_float_type)
+                face_vertices = np.zeros((nvf, 3), dtype=self.np_float_type)
+                self.finalize_extra_verts(self.sdf_AF(esdf_l), self.sdf_AF(esdf_r), self.AF(edge_vertices), self.sdf_AF(fsdf_l), self.sdf_AF(fsdf_r), self.AF(face_vertices))
+                del esdf_l, esdf_r, fsdf_l, fsdf_r
+                faces = AC(np.zeros((nf, 3), dtype=np.int32))
+                self.get_faces(AsInt(faces))
+                vertices = np.concatenate((vertices, edge_vertices, face_vertices))
+                in_view_tag = np.zeros(vertices.shape[0], dtype=bool)
+                self.get_in_view_tag(e, AsBool(in_view_tag))
+                in_view_tags.append(in_view_tag)
+                meshes.append(trimesh.Trimesh(vertices=vertices, faces=faces, process=False))
+                print(f"element {e} has vertices #{meshes[-1].vertices.shape[0]} faces #{meshes[-1].faces.shape[0]}")
+        return meshes, in_view_tags