import gradio as gr
import trimesh
from PIL import Image
import base64
import io

def modify_texture():
    # Load the GLB file
    scene = trimesh.load_mesh("train.glb")
    geometries = list(scene.geometry.values())
    geometry = geometries[0]
    adjacency_matrix = geometry.edges_sparse
    
    # Open the texture image
    im = Image.open("defect.jpg")
    
    # Create a material with the texture image
    material = trimesh.visual.texture.SimpleMaterial(image=im)
    
    # Get the uv coordinates from the mesh
    uv = adjacency_matrix.visual.uv
    
    # Create TextureVisuals with the uv coordinates and the texture image
    color_visuals = trimesh.visual.TextureVisuals(uv=uv, image=im, material=material)
    
    # Apply the visuals to the mesh
    mesh.visual = color_visuals
    
    # Export the mesh to a .glb file
    mesh.export(file_obj='modified_train.glb')
    
    # Convert the GLB file to base64 for Gradio Model3D output
    with open('modified_train.glb', 'rb') as f:
        glb_base64 = base64.b64encode(f.read()).decode()
    
    return glb_base64

gr.Interface(
    fn=modify_texture,
    inputs=[],
    outputs=gr.Model3D(label="Download Modified GLB File"),
    title="GLB Texture Modifier"
).launch()