import gradio as gr
import numpy as np
from PIL import Image
import trimesh
import tempfile
import os

def image_to_3d(image, scale=1.0):
    # Convert image to grayscale
    img = Image.fromarray(image).convert('L')
    img_array = np.array(img)
    
    # Normalize image values
    height_data = img_array / 255.0
    
    # Create grid
    x, y = np.mgrid[0:height_data.shape[0], 0:height_data.shape[1]]
    
    # Create vertices
    vertices = np.column_stack([x.flatten(), y.flatten(), height_data.flatten() * scale])
    
    # Create faces
    faces = []
    for i in range(height_data.shape[0]-1):
        for j in range(height_data.shape[1]-1):
            v0 = i * height_data.shape[1] + j
            v1 = v0 + 1
            v2 = v0 + height_data.shape[1]
            v3 = v2 + 1
            faces.append([v0, v1, v2])
            faces.append([v1, v3, v2])
    
    # Create mesh
    mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
    
    # Save to temporary STL file
    with tempfile.NamedTemporaryFile(suffix=".stl", delete=False) as tmp:
        mesh.export(tmp.name, file_type='stl')
        return tmp.name

def process_image(image, scale_factor):
    scale = float(scale_factor)
    stl_path = image_to_3d(image, scale)
    return stl_path, gr.File.update(value=stl_path, visible=True)

with gr.Blocks() as demo:
    gr.Markdown("# Image to 3D Model Converter")
    gr.Markdown("Upload an image to generate a 3D printable STL file")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload Image")
            scale_slider = gr.Slider(1, 10, value=1.0, label="Scale Factor")
            generate_btn = gr.Button("Generate 3D Model")
        
        with gr.Column():
            model_view = gr.Model3D(label="3D Model Preview", clear_color=[0.8, 0.8, 0.8, 1.0])
            download_output = gr.File(label="Download STL File", visible=False)
    
    generate_btn.click(
        fn=process_image,
        inputs=[image_input, scale_slider],
        outputs=[model_view, download_output]
    )

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