import streamlit as st
import cv2
import numpy as np
from PIL import Image
import os
from image_processor import ImageProcessor
from model_generator import ModelGenerator
from utils import create_output_directory, allowed_file

def main():
    st.title("Image to 3D Converter")
    st.write("Upload an image to convert it to a 3D model")

    # Initialize processors
    image_processor = ImageProcessor()
    model_generator = ModelGenerator()

    # File uploader
    uploaded_file = st.file_uploader("Choose an image file", type=['png', 'jpg', 'jpeg'])

    if uploaded_file is not None:
        if allowed_file(uploaded_file.name):
            # Read and display the uploaded image
            image = Image.open(uploaded_file)
            st.image(image, caption="Uploaded Image", use_column_width=True)

            # Convert PIL Image to numpy array
            image_np = np.array(image)

            # Process button
            if st.button("Convert to 3D"):
                with st.spinner("Processing..."):
                    # Generate depth map
                    depth_map = image_processor.generate_depth_map(image_np)
                    
                    # Display depth map
                    st.image(depth_map, caption="Generated Depth Map", use_column_width=True)

                    # Generate 3D model
                    points = model_generator.depth_to_points(depth_map, scale=50.0)
                    mesh = model_generator.generate_mesh(points)

                    # Create output directory
                    output_dir = create_output_directory()
                    output_file = os.path.join(output_dir, "output_model.obj")

                    # Save the mesh
                    model_generator.save_mesh(mesh, output_file)
                    
                    # Provide download link
                    with open(output_file, "rb") as file:
                        st.download_button(
                            label="Download 3D Model",
                            data=file,
                            file_name="3d_model.obj",
                            mime="application/octet-stream"
                        )

                st.success("3D model generated successfully!")
        else:
            st.error("Please upload an image file (PNG, JPG, or JPEG)")

if __name__ == "__main__":
    main()