Update app.py

app.py

@@ -4,11 +4,9 @@ import numpy as np
 from PIL import Image
 import tempfile
 
-def apply_masked_texture(section):
+def apply_texture_with_mask(section):
     # Load the original mesh
     mesh = trimesh.load('train.glb', force='mesh')
-
-    # Load the new texture image
     im = Image.open('rust_steel.png').convert('RGB')
 
     # Calculate bounding box coordinates
@@ -18,7 +16,7 @@ def apply_masked_texture(section):
     mid_y = (max_bounds[1] + min_bounds[1]) / 2
     mid_z = (max_bounds[2] + min_bounds[2]) / 2
 
-    # Define sections
+    # Define sections with indices
     sections = {
         'upper': np.where(mesh.vertices[:, 2] > mid_z)[0],
         'lower': np.where(mesh.vertices[:, 2] <= mid_z)[0],
@@ -29,23 +27,27 @@ def apply_masked_texture(section):
         'right': np.where(mesh.vertices[:, 0] > mid_x)[0]
     }
 
-    # Get indices for the selected section
+    # Select the indices for the specified section
     selected_indices = sections[section]
 
-    # Copy the original visual properties (colors, UVs)
+    # Preserve the original visual appearance
     if mesh.visual.vertex_colors is not None:
-        new_vertex_colors = np.copy(mesh.visual.vertex_colors)
+        vertex_colors = np.copy(mesh.visual.vertex_colors)
     else:
-        new_vertex_colors = np.ones((len(mesh.vertices), 4)) * 255  # Assuming original is white or neutral
+        vertex_colors = np.ones((len(mesh.vertices), 4)) * 255  # Assuming the default color is white
+
+    # Apply the new texture only to the selected section
+    uv = np.random.rand(len(mesh.vertices), 2)  # Generate new UV coordinates
+    uv_colors = trimesh.visual.uv_to_color(uv[selected_indices], im)  # Map UV coordinates to colors in the image
 
-    # Apply the new texture to only the selected section
-    for idx in selected_indices:
-        new_vertex_colors[idx, :] = [255, 255, 0, 255]  # Example yellow RGBA color
+    for i, idx in enumerate(selected_indices):
+        vertex_colors[idx] = uv_colors[i]  # Apply the new texture color to selected indices
 
-    # Apply the new visual properties
-    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=trimesh.visual.ColorVisuals(vertex_colors=new_vertex_colors), validate=True, process=False)
+    # Apply the new visual colors to the mesh
+    color_visuals = trimesh.visual.ColorVisuals(vertex_colors=vertex_colors)
+    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=color_visuals, validate=True, process=False)
 
-    # Save to a temporary file for Gradio visualization
+    # Save the modified mesh to a temporary file for Gradio visualization
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.glb')
     textured_mesh.export(temp_file.name, file_type='glb')
     temp_file.close()
@@ -58,6 +60,6 @@ with gr.Blocks() as app:
 
     section_dropdown = gr.Dropdown(choices=['upper', 'lower', 'middle', 'top', 'right'], label="Select Section")
     button = gr.Button("Apply Texture to Section")
-    button.click(apply_masked_texture, inputs=section_dropdown, outputs=modified_model)
+    button.click(apply_texture_with_mask, inputs=section_dropdown, outputs=modified_model)
 
 app.launch()