import gradio as gr
import trimesh
from instant_texture import Converter

converter = Converter()


def extract_texture(mesh_path: str):
    mesh = trimesh.load(mesh_path)
    mesh = next(iter(mesh.geometry.values()))
    texture = mesh.visual.material.baseColorTexture
    return texture


def convert(input_mesh_path: str):
    if not input_mesh_path.endswith(".obj"):
        raise ValueError("Only .obj files are supported")
    output_path = "/tmp/output.glb"
    converter.convert(input_mesh_path, output_path)
    texture = extract_texture(output_path)
    return output_path, texture


demo = gr.Interface(
    convert,
    title="Instant Texture",
    description="Convert a vertex-colored mesh (.obj) to a uv-mapped, textured mesh (.glb).",
    inputs=gr.Model3D(label="Vertex-colored mesh (.obj)"),
    outputs=[
        gr.Model3D(label="Output uv-mapped, textured mesh (.glb)"),
        gr.Image(label="Output texture"),
    ],
    examples=["examples/chair.obj"],
    cache_examples=True,
)

if __name__ == "__main__":
    demo.queue().launch()