Code Editor for app.py

import streamlit as st
import numpy as np
import cv2
import plotly.graph_objects as go
from scipy import ndimage
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)
    
    # Get room boundaries
    rooms = []
    for contour in contours:
        if cv2.contourArea(contour) > 100:  # Filter small contours
            x, y, w, h = cv2.boundingRect(contour)
            rooms.append({
                'x': x,
                'y': y,
                'width': w,
                'height': h
            })
    
    return rooms

def create_3d_visualization(rooms):
    # Create 3D visualization using Plotly
    fig = go.Figure()
    
    # Standard room height
    height = 3
    
    for room in rooms:
        x = [room['x'], room['x'], room['x']+room['width'], room['x']+room['width'], room['x']]
        y = [room['y'], room['y'], room['y'], room['y'], room['y']]
        z = [0, height, height, 0, 0]
        
        # Add walls
        fig.add_trace(go.Scatter3d(
            x=x, y=y, z=z,
            mode='lines',
            line=dict(color='blue', width=2),
            showlegend=False
        ))
        
        # Add floor
        fig.add_trace(go.Mesh3d(
            x=[room['x'], room['x']+room['width'], room['x']+room['width'], room['x']],
            y=[room['y'], room['y'], room['y']+room['height'], room['y']+room['height']],
            z=[0, 0, 0, 0],
            color='lightgray',
            opacity=0.7,
            showlegend=False
        ))
        
        # Add label
        fig.add_trace(go.Scatter3d(
            x=[room['x'] + room['width']/2],
            y=[room['y'] + room['height']/2],
            z=[height/2],
            mode='text',
            text=['Room'],
            textposition='middle center',
            showlegend=False
        ))
    
    fig.update_layout(
        scene=dict(
            xaxis_title='X',
            yaxis_title='Y',
            zaxis_title='Height',
            aspectmode='data'
        ),
        margin=dict(l=0, r=0, t=0, b=0)
    )
    
    return fig

st.write("Upload a 2D floor plan image to convert it into a 3D visualization.")
st.write("Note: The floor plan should have clear walls and room boundaries.")

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

if uploaded_file is not None:
    # Read 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 floor plan
    rooms = process_floor_plan(image)
    
    if rooms:
        st.subheader("3D Visualization")
        fig = create_3d_visualization(rooms)
        st.plotly_chart(fig, use_container_width=True)
        
        st.write("Instructions:")
        st.write("- Use mouse to rotate the 3D view")
        st.write("- Scroll to zoom in/out")
        st.write("- Double click to reset view")
    else:
        st.error("No rooms detected in the floor plan. Please ensure the image has clear boundaries.")