Upload app.py

app.py

@@ -1,7 +1,88 @@
 import multiprocessing
 import objaverse
 import random
-import numpy as np
+import numpy
+import trimesh
+
+
+
+def model_to_pc(mesh: trimesh.Trimesh, n_sample_points=10000):
+    f32 = numpy.float32
+    rad = numpy.sqrt(mesh.area / (3 * n_sample_points))
+    for _ in range(24):
+        pcd, face_idx = trimesh.sample.sample_surface_even(mesh, n_sample_points, rad)
+        rad *= 0.85
+        if len(pcd) == n_sample_points:
+            break
+    else:
+        raise ValueError("Bad geometry, cannot finish sampling.", mesh.area)
+    if isinstance(mesh.visual, trimesh.visual.ColorVisuals):
+        rgba = mesh.visual.face_colors[face_idx]
+    elif isinstance(mesh.visual, trimesh.visual.TextureVisuals):
+        bc = trimesh.proximity.points_to_barycentric(mesh.triangles[face_idx], pcd)
+        if mesh.visual.uv is None or len(mesh.visual.uv) < mesh.faces[face_idx].max():
+            uv = numpy.zeros([len(bc), 2])
+            # st.warning("Invalid UV, filling with zeroes")
+        else:
+            uv = numpy.einsum('ntc,nt->nc', mesh.visual.uv[mesh.faces[face_idx]], bc)
+        material = mesh.visual.material
+        if hasattr(material, 'materials'):
+            if len(material.materials) == 0:
+                rgba = numpy.ones_like(pcd) * 0.8
+                texture = None
+                # st.warning("Empty MultiMaterial found, falling back to light grey")
+            else:
+                material = material.materials[0]
+        if hasattr(material, 'image'):
+            texture = material.image
+            if texture is None:
+                rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color
+        elif hasattr(material, 'baseColorTexture'):
+            texture = material.baseColorTexture
+            if texture is None:
+                rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color
+        else:
+            texture = None
+            rgba = numpy.ones_like(pcd) * 0.8
+            # st.warning("Unknown material, falling back to light grey")
+        if texture is not None:
+            rgba = trimesh.visual.uv_to_interpolated_color(uv, texture)
+    if rgba.max() > 1:
+        if rgba.max() > 255:
+            rgba = rgba.astype(f32) / rgba.max()
+        else:
+            rgba = rgba.astype(f32) / 255.0
+    return numpy.concatenate([numpy.array(pcd, f32), numpy.array(rgba, f32)[:, :3]], axis=-1)
+
+
+def trimesh_to_pc(scene_or_mesh):
+    if isinstance(scene_or_mesh, trimesh.Scene):
+        meshes = []
+        for node_name in scene_or_mesh.graph.nodes_geometry:
+            # which geometry does this node refer to
+            transform, geometry_name = scene_or_mesh.graph[node_name]
+
+            # get the actual potential mesh instance
+            geometry = scene_or_mesh.geometry[geometry_name].copy()
+            if not hasattr(geometry, 'triangles'):
+                continue
+            geometry: trimesh.Trimesh
+            geometry = geometry.apply_transform(transform)
+            meshes.append(geometry)
+        total_area = sum(geometry.area for geometry in meshes)
+        if total_area < 1e-6:
+            raise ValueError("Bad geometry: total area too small (< 1e-6)")
+        pcs = []
+        for geometry in meshes:
+            pcs.append(model_to_pc(geometry, max(1, round(geometry.area / total_area * 10000))))
+        if not len(pcs):
+            raise ValueError("Unsupported mesh object: no triangles found")
+        return numpy.concatenate(pcs)
+    else:
+        assert isinstance(scene_or_mesh, trimesh.Trimesh)
+        return model_to_pc(scene_or_mesh, 10000)
+
+
 
 processes = multiprocessing.cpu_count()
 
@@ -11,4 +92,17 @@ random_object_uids = random.sample(uids, 100)
 objects = objaverse.load_objects(
     uids=random_object_uids,
     download_processes=processes
-)
+)
+
+pc_list = []
+for objaid in random_object_uids:
+    objamodel = objaverse.load_objects([objaid])[objaid]
+    pc = trimesh_to_pc(trimesh.load(objamodel))
+    pc_list.append(pc.reshape([1, -1, pc.shape[-1]]))
+
+pc_list = numpy.concatenate(pc_list, axis=0)
+
+numpy.save("pc.npy", pc_list)
+
+print("Done")
+