Fix loading and render.

app.py

@@ -58,7 +58,8 @@ def load_data():
 
 
 def render_pc(pc):
-    pc = pc[:2048]
+    rand = numpy.random.permutation(len(pc))[:2048]
+    pc = pc[rand]
     rgb = (pc[:, 3:] * 255).astype(numpy.uint8)
     g = go.Scatter3d(
         x=pc[:, 0], y=pc[:, 1], z=pc[:, 2],
@@ -66,6 +67,7 @@ def render_pc(pc):
         marker=dict(size=2, color=[f'rgb({rgb[i, 0]}, {rgb[i, 1]}, {rgb[i, 2]})' for i in range(len(pc))]),
     )
     fig = go.Figure(data=[g])
+    fig.update_layout(scene_camera=dict(up=dict(x=0, y=1, z=0)))
     st.plotly_chart(fig)
     # st.caption("Point Cloud Preview")
 
@@ -87,10 +89,10 @@ try:
     with tab_pc2img:
         prompt = st.text_input("Prompt")
         noise_scale = st.slider('Variation Level', 0, 5)
-        cfg_scale = st.slider('Guidance Scale', 0.0, 30.0, 10.0)
-        steps = st.slider('Diffusion Steps', 2, 80, 10)
-        width = st.slider('Width', 128, 512, step=32)
-        height = st.slider('Height', 128, 512, step=32)
+        cfg_scale = st.slider('Guidance Scale', 0.0, 30.0, 3.0)
+        steps = st.slider('Diffusion Steps', 8, 80, 10)
+        width = st.slider('Width', 480, 640, step=32)
+        height = st.slider('Height', 480, 640, step=32)
         if st.button("Generate"):
             pc = load_data()
             render_pc(pc)

misc_utils.py

@@ -34,7 +34,24 @@ def model_to_pc(mesh: trimesh.Trimesh, n_sample_points=10000):
     elif isinstance(mesh.visual, trimesh.visual.TextureVisuals):
         bc = trimesh.proximity.points_to_barycentric(mesh.triangles[face_idx], pcd)
         uv = numpy.einsum('ntc,nt->nc', mesh.visual.uv[mesh.faces[face_idx]], bc)
-        rgba = trimesh.visual.uv_to_interpolated_color(uv, mesh.visual.material.image)
+        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
+        elif hasattr(material, 'baseColorTexture'):
+            texture = material.baseColorTexture
+        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()
@@ -45,11 +62,18 @@ def model_to_pc(mesh: trimesh.Trimesh, n_sample_points=10000):
 
 def trimesh_to_pc(scene_or_mesh):
     if isinstance(scene_or_mesh, trimesh.Scene):
-        meshes = [
-            model_to_pc(trimesh.Trimesh(vertices=g.vertices, faces=g.faces), 10000 // len(scene_or_mesh.geometry))
-            for g in scene_or_mesh.geometry.values()
-            if isinstance(g, trimesh.Trimesh)
-        ]
+        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(model_to_pc(geometry, 10000 // len(scene_or_mesh.geometry)))
         if not len(meshes):
             return None
         return numpy.concatenate(meshes)