Update app.py

app.py

@@ -287,6 +287,17 @@ def process_image(image, model_choice="GLPN (Recommended)", visualization_type="
             pcd, depth=10, n_threads=1
         )[0]
         
+        # Transfer colors from point cloud to mesh vertices
+        print("Transferring colors to mesh...")
+        pcd_tree = o3d.geometry.KDTreeFlann(pcd)
+        mesh_colors = []
+        for vertex in mesh.vertices:
+            # Find nearest point in point cloud
+            [_, idx, _] = pcd_tree.search_knn_vector_3d(vertex, 1)
+            # Get color from nearest point
+            mesh_colors.append(pcd.colors[idx[0]])
+        mesh.vertex_colors = o3d.utility.Vector3dVector(np.array(mesh_colors))
+        
         # Rotate mesh
         rotation = mesh.get_rotation_matrix_from_xyz((np.pi, 0, 0))
         mesh.rotate(rotation, center=(0, 0, 0))
@@ -345,13 +356,31 @@ def process_image(image, model_choice="GLPN (Recommended)", visualization_type="
             vertices = np.asarray(mesh.vertices)
             triangles = np.asarray(mesh.triangles)
             
-            mesh_trace = go.Mesh3d(
-                x=vertices[:, 0], y=vertices[:, 1], z=vertices[:, 2],
-                i=triangles[:, 0], j=triangles[:, 1], k=triangles[:, 2],
-                color='lightblue',
-                opacity=0.9,
-                name='Mesh'
-            )
+            # Get vertex colors if available
+            if mesh.has_vertex_colors():
+                vertex_colors = np.asarray(mesh.vertex_colors)
+                colors_rgb = ['rgb({},{},{})'.format(int(r*255), int(g*255), int(b*255)) 
+                              for r, g, b in vertex_colors]
+                
+                mesh_trace = go.Mesh3d(
+                    x=vertices[:, 0], y=vertices[:, 1], z=vertices[:, 2],
+                    i=triangles[:, 0], j=triangles[:, 1], k=triangles[:, 2],
+                    vertexcolor=colors_rgb,
+                    opacity=0.95,
+                    name='Mesh',
+                    lighting=dict(ambient=0.5, diffuse=0.8, specular=0.2),
+                    lightposition=dict(x=100, y=100, z=100)
+                )
+            else:
+                # Fallback if no colors
+                mesh_trace = go.Mesh3d(
+                    x=vertices[:, 0], y=vertices[:, 1], z=vertices[:, 2],
+                    i=triangles[:, 0], j=triangles[:, 1], k=triangles[:, 2],
+                    color='lightblue',
+                    opacity=0.9,
+                    name='Mesh'
+                )
+            
             data.append(mesh_trace)
         
         layout = go.Layout(