Update app.py

app.py

@@ -4,19 +4,19 @@ import numpy as np
 from PIL import Image
 import tempfile
 
-def apply_masked_texture(section):
+def apply_masked_visuals(section):
     # Load the original mesh
     mesh = trimesh.load('train.glb', force='mesh')
     im = Image.open('rust_steel.png').convert('RGB')
 
-    # Calculate the bounding box coordinates
+    # Calculate bounding box coordinates
     bounds = mesh.bounds
     min_bounds, max_bounds = bounds[0], bounds[1]
     mid_x = (max_bounds[0] + min_bounds[0]) / 2
     mid_y = (max_bounds[1] + min_bounds[1]) / 2
     mid_z = (max_bounds[2] + min_bounds[2]) / 2
 
-    # Define sections with relevant indices
+    # Define sections with indices
     sections = {
         'upper': np.where(mesh.vertices[:, 2] > mid_z)[0],
         'lower': np.where(mesh.vertices[:, 2] <= mid_z)[0],
@@ -30,26 +30,24 @@ def apply_masked_texture(section):
     # Select the indices for the specified section
     selected_indices = sections[section]
 
-    # Initialize UV mapping and retain original visuals for unselected vertices
-    if mesh.visual.kind == 'texture' and mesh.visual.uv is not None:
-        original_uv = mesh.visual.uv
+    # Check if the mesh has existing visuals
+    if mesh.visual.kind == 'texture':
+        visual = mesh.visual.copy()
+        visual.update_vertices(mask=np.isin(range(len(mesh.vertices)), selected_indices))
     else:
-        original_uv = np.random.rand(len(mesh.vertices), 2)
+        visual = trimesh.visual.ColorVisuals(vertex_colors=np.ones((len(mesh.vertices), 4)) * 255)
 
-    # Create new UV coordinates only for the selected region
-    new_uv = np.copy(original_uv)
+    # Prepare the new visual properties (UV or color updates)
+    uv = np.copy(visual.uv) if visual.kind == 'texture' else np.random.rand(len(mesh.vertices), 2)
+    new_uv = np.copy(uv)
     new_uv[selected_indices, :] = np.random.rand(len(selected_indices), 2)
 
-    # Preserve the original colors and visuals for non-selected sections
-    original_colors = mesh.visual.vertex_colors if mesh.visual.vertex_colors is not None else np.ones((len(mesh.vertices), 4)) * 255
-
-    # Create a new material with the texture image
+    # Create a new material with the image texture
     material = trimesh.visual.texture.SimpleMaterial(image=im)
     new_visuals = trimesh.visual.TextureVisuals(uv=new_uv, material=material)
 
-    # Create a new mesh with both original and updated visuals
+    # Create a new mesh with the updated visuals
     textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=new_visuals, validate=True, process=False)
-    textured_mesh.visual.vertex_colors = original_colors  # Retain original vertex colors
 
     # Save the modified mesh to a temporary file for Gradio visualization
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.glb')
@@ -64,6 +62,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_masked_visuals, inputs=section_dropdown, outputs=modified_model)
 
 app.launch()