import gradio as gr
import trimesh
import numpy as np
from PIL import Image
import io

def visualize_texture():
    # Load the GLB file
    mesh = trimesh.load('train.glb', file_type='glb', process=False)

    # Load the texture image
    im = Image.open('defect.jpg')
    im = im.convert('RGB')  # Ensure the image is in RGB format

    # Create random UV coordinates for the mesh
    uv = np.random.rand(len(mesh.vertices), 2)

    # Create material and apply texture
    material = trimesh.visual.texture.SimpleMaterial(image=im)
    color_visuals = trimesh.visual.TextureVisuals(uv=uv, image=im, material=material)

    # Apply the texture to the original mesh
    textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=color_visuals, validate=True, process=False)

    # Save the original and textured meshes to GLB format in memory
    original_glb = io.BytesIO()
    mesh.export(original_glb, file_type='glb')
    original_glb.seek(0)

    modified_glb = io.BytesIO()
    textured_mesh.export(modified_glb, file_type='glb')
    modified_glb.seek(0)

    return original_glb, modified_glb

with gr.Blocks() as app:
    gr.Markdown("### 3D Model Texture Application")
    original_model = gr.Model3D(label="Original Model")
    modified_model = gr.Model3D(label="Textured Model")
    button = gr.Button("Visualize Texture")
    button.click(visualize_texture, outputs=[original_model, modified_model])

app.launch()