Code Editor for app.py

import streamlit as st
import numpy as np
import cv2
import plotly.graph_objects as go
from skimage import measure
import io
from PIL import Image

st.title("2D to 3D Floor Plan Converter")

def process_floor_plan(image):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Threshold the image
    _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY_INV)
    
    # Find contours
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Create wall coordinates
    walls = []
    for contour in contours:
        for point in contour:
            x, y = point[0]
            walls.append([x, y, 0])  # Ground level
            walls.append([x, y, 30])  # Wall height
    
    return np.array(walls)

def create_3d_visualization(walls):
    # Create 3D scatter plot
    fig = go.Figure(data=[
        go.Scatter3d(
            x=walls[:, 0],
            y=walls[:, 1],
            z=walls[:, 2],
            mode='markers',
            marker=dict(
                size=2,
                color='blue',
                opacity=0.8
            )
        )
    ])
    
    # Update layout
    fig.update_layout(
        scene=dict(
            xaxis_title='Width',
            yaxis_title='Length',
            zaxis_title='Height',
            aspectmode='data'
        ),
        width=800,
        height=800,
        showlegend=False
    )
    
    return fig

st.write("Upload a 2D floor plan image (black and white preferred)")

uploaded_file = st.file_uploader("Choose a floor plan image", type=['png', 'jpg', 'jpeg'])

if uploaded_file is not None:
    # Read and process the image
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    
    # Display original image
    st.subheader("Original Floor Plan")
    st.image(uploaded_file, use_column_width=True)
    
    # Process the floor plan
    try:
        walls = process_floor_plan(image)
        
        # Create and display 3D visualization
        st.subheader("3D Visualization")
        st.write("Use your mouse to rotate, zoom, and pan the 3D model")
        fig = create_3d_visualization(walls)
        st.plotly_chart(fig)
        
        st.info("Note: This is a simplified 3D representation. The walls are represented as points in 3D space. For more accurate results, consider using specialized architectural software.")
        
    except Exception as e:
        st.error("Error processing the floor plan. Please ensure the image is clear and has good contrast between walls and spaces.")
        st.error(f"Error details: {str(e)}")