Add a function to sample points non-uniformly

mesh_to_sdf/__init__.py

@@ -1,6 +1,6 @@
 import numpy as np
 import mesh_to_sdf.surface_point_cloud
-from mesh_to_sdf.utils import scale_to_unit_cube, get_raster_points, check_voxels
+from mesh_to_sdf.utils import scale_to_unit_cube, scale_to_unit_sphere, get_raster_points, check_voxels
 import trimesh
 
 class BadMeshException(Exception):
@@ -56,4 +56,39 @@ def mesh_to_voxels(mesh, voxel_resolution=64, surface_point_method='scan', sign_
     if pad:
         voxels = np.pad(voxels, 1, mode='constant', constant_values=1)
 
-    return voxels
+    return voxels
+
+# Sample some uniform points and some normally distributed around the surface as proposed in the DeepSDF paper
+def sample_sdf_near_surface(mesh, number_of_points = 500000, surface_point_method='scan', sign_method='normal', scan_count=100, scan_resolution=400, sample_point_count=10000000, normal_sample_count=11, min_size=0):
+    mesh = scale_to_unit_sphere(mesh)
+
+    if surface_point_method == 'sample' and sign_method == 'depth':
+        print("Incompatible methods for sampling points and determining sign, using sign_method='normal' instead.")
+        sign_method = 'normal'
+
+    surface_point_cloud = get_surface_point_cloud(mesh, surface_point_method, 1, scan_count, scan_resolution, sample_point_count, calculate_normals=sign_method=='normal')
+
+    query_points = []
+    surface_sample_count = int(number_of_points * 47 / 50) // 2
+    surface_points = surface_point_cloud.get_random_surface_points(surface_sample_count, use_scans=surface_point_method=='scan')
+    query_points.append(surface_points + np.random.normal(scale=0.0025, size=(surface_sample_count, 3)))
+    query_points.append(surface_points + np.random.normal(scale=0.00025, size=(surface_sample_count, 3)))
+
+    unit_sphere_sample_count = number_of_points - surface_sample_count
+    unit_sphere_points = np.random.uniform(-1, 1, size=(unit_sphere_sample_count * 2, 3))
+    unit_sphere_points = unit_sphere_points[np.linalg.norm(unit_sphere_points, axis=1) < 1]
+    query_points.append(unit_sphere_points[:unit_sphere_sample_count, :])
+    query_points = np.concatenate(query_points).astype(np.float32)
+
+    if sign_method == 'normal':
+        sdf = surface_point_cloud.get_sdf_in_batches(query_points, use_depth_buffer=False)
+    elif sign_method == 'depth':
+        sdf = surface_point_cloud.get_sdf_in_batches(query_points, use_depth_buffer=True, sample_count=sample_point_count)
+    else:
+        raise ValueError('Unknown sign determination method: {:s}'.format(sign_method))
+
+    model_size = np.count_nonzero(sdf[-unit_sphere_sample_count:] < 0) / unit_sphere_sample_count
+    if model_size < min_size:
+        raise BadMeshException()
+
+    return query_points, sdf

mesh_to_sdf/surface_point_cloud.py

@@ -52,29 +52,6 @@ def get_sdf_in_batches(self, query_points, use_depth_buffer=False, sample_count=
             end = min(result.shape[0], (i + 1) * batch_size)
             result[start:end] = self.get_sdf(query_points[start:end, :], use_depth_buffer=use_depth_buffer, sample_count=sample_count)
         return result
-
-    def get_sample_points(self, number_of_points = 200000):
-        ''' Use sample points as described in the DeepSDF paper '''
-        points = []
-
-        surface_sample_count = int(number_of_points * 0.4)
-        surface_points = self.get_random_surface_points(surface_sample_count)
-        points.append(surface_points + np.random.normal(scale=0.0025, size=(surface_sample_count, 3)))
-        points.append(surface_points + np.random.normal(scale=0.00025, size=(surface_sample_count, 3)))
-
-        unit_sphere_sample_count = int(number_of_points * 0.2)
-        unit_sphere_points = np.random.uniform(-1, 1, size=(unit_sphere_sample_count * 2, 3))
-        unit_sphere_points = unit_sphere_points[np.linalg.norm(unit_sphere_points, axis=1) < 1]
-        points.append(unit_sphere_points[:unit_sphere_sample_count, :])
-        points = np.concatenate(points).astype(np.float32)
-
-        sdf = self.get_sdf(points)
-
-        model_size = np.count_nonzero(sdf[-unit_sphere_sample_count:] < 0) / unit_sphere_sample_count
-        if model_size < 0.015:
-            raise BadMeshException()
-
-        return points, sdf
 
     def get_surface_points_and_normals(self, number_of_points = 50000):
         count = self.points.shape[0]