Update app.py

app.py

@@ -3,60 +3,56 @@ import trimesh
 import numpy as np
 from PIL import Image
 import tempfile
+import os
 
-def apply_rust_texture(section):
-    # Load the original mesh
+def visualize_texture(section):
     mesh = trimesh.load('train.glb', force='mesh')
-    rust_texture = Image.open('rust_steel.png').convert('RGB')
+    im = Image.open('rust_steel.png').convert('RGB')
 
-    # Train body selection logic
+    # Calculate bounds
     bounds = mesh.bounds
     min_bounds, max_bounds = bounds[0], bounds[1]
-    body_height = max_bounds[2] - min_bounds[2]
-    body_midpoint = (max_bounds + min_bounds) / 2
-
-    # Define more sophisticated criteria for the body selection
-    if section == 'body':
-        # Assuming the body is within certain height and centered horizontally
-        selected_indices = np.where((mesh.vertices[:, 2] >= min_bounds[2]) & 
-                                    (mesh.vertices[:, 2] <= min_bounds[2] + 0.75 * body_height) &
-                                    (mesh.vertices[:, 0] >= min_bounds[0]) &
-                                    (mesh.vertices[:, 0] <= max_bounds[0]))[0]
-
-    elif section == 'roof':
-        # Roof is the top quarter of the height
-        selected_indices = np.where(mesh.vertices[:, 2] > min_bounds[2] + 0.75 * body_height)[0]
-
-    # Ensure indices are within bounds
-    valid_indices = selected_indices[selected_indices < len(mesh.vertices)]
-
-    # Create or initialize UV mapping
+    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
+    sections = {
+        'upper': np.where(mesh.vertices[:, 2] > mid_z)[0],
+        'lower': np.where(mesh.vertices[:, 2] <= mid_z)[0],
+        'middle': np.where((mesh.vertices[:, 0] >= min_bounds[0]) & (mesh.vertices[:, 0] <= max_bounds[0]) &
+                           (mesh.vertices[:, 1] >= min_bounds[1]) & (mesh.vertices[:, 1] <= max_bounds[1]) &
+                           (mesh.vertices[:, 2] >= min_bounds[2]) & (mesh.vertices[:, 2] <= max_bounds[2]))[0],
+        'top': np.where(mesh.vertices[:, 1] > mid_y)[0],
+        'right': np.where(mesh.vertices[:, 0] > mid_x)[0]
+    }
+
+    # Get indices for selected section
+    selected_indices = sections[section]
+
+    # Create UV coordinates
     uv = np.random.rand(len(mesh.vertices), 2)
+    new_uv = np.zeros_like(uv)
+    new_uv[selected_indices, :] = uv[selected_indices, :]
 
-    # Apply the rust texture only to the selected indices
-    new_uv = np.copy(uv)
-    new_uv[valid_indices, :] = np.random.rand(len(valid_indices), 2)
+    # Create material and apply texture
+    material = trimesh.visual.texture.SimpleMaterial(image=im)
+    color_visuals = trimesh.visual.TextureVisuals(uv=new_uv, image=im, material=material)
+    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=color_visuals, validate=True, process=False)
 
-    # Create a new material with the rust texture image
-    material = trimesh.visual.texture.SimpleMaterial(image=rust_texture)
-    new_visuals = trimesh.visual.TextureVisuals(uv=new_uv, material=material)
-
-    # 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)
-
-    # Save the modified mesh to a temporary file for Gradio visualization
+    # Save the mesh to a temporary file
     temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.glb')
     textured_mesh.export(temp_file.name, file_type='glb')
-    temp_file.close()
+    temp_file.close()  # Close the file so it can be read by Gradio
     return temp_file.name
 
 with gr.Blocks() as app:
-    gr.Markdown("### Apply Rust Texture to Train Body Section")
+    gr.Markdown("### 3D Model Texture Application")
     original_model = gr.Model3D('train.glb', label="Original Model")
     modified_model = gr.Model3D(label="Textured Model")
-
-    section_dropdown = gr.Dropdown(choices=['body', 'roof'], label="Select Section")
-    button = gr.Button("Apply Rust Texture to Section")
-    button.click(apply_rust_texture, inputs=section_dropdown, outputs=modified_model)
+    
+    section_dropdown = gr.Dropdown(choices=['upper', 'lower', 'middle', 'top', 'right'], label="Select Section")
+    button = gr.Button("Visualize Texture")
+    button.click(visualize_texture, inputs=section_dropdown, outputs=modified_model)
 
 app.launch()