import trimesh
import gradio as gr
import os

def generate_3d_object(shape_type, radius, width, height):
    if shape_type == "Sphere":
        mesh = trimesh.creation.icosphere(radius=radius)
    elif shape_type == "Box":
        mesh = trimesh.creation.box(extents=[width, height, radius])
    else:
        raise ValueError("Unsupported shape type")
    
    # Save the mesh to a file
    output_file = "generated_object.obj"
    mesh.export(output_file)
    
    return output_file

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_3d_object,
    inputs=[
        gr.Dropdown(choices=["Sphere", "Box"], label="Shape Type"),
        gr.Number(label="Radius"),
        gr.Number(label="Width (for Box)"),
        gr.Number(label="Height (for Box)")
    ],
    outputs=gr.Model3D(label="Generated 3D Object"),
    title="3D Object Generator",
    description="Generate 3D objects like spheres and boxes.",
    flagging_dir=os.path.join(os.getcwd(), "flagged")  # Specify the flagging directory
)

# Launch the interface
iface.launch()