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 base64
from io import BytesIO
from PIL import Image

st.set_page_config(layout="wide")
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)
    
    # Process contours to get room information
    rooms = []
    for contour in contours:
        if cv2.contourArea(contour) > 100:  # Filter small contours
            # Get bounding box
            x, y, w, h = cv2.boundingRect(contour)
            center = (x + w//2, y + h//2)
            rooms.append({
                'contour': contour,
                'center': center,
                'width': w,
                'height': h,
                'area': cv2.contourArea(contour)
            })
    
    return rooms

def create_3d_visualization(rooms):
    # Create 3D visualization using Plotly
    fig = go.Figure()
    
    height = 3  # Standard room height
    
    for i, room in enumerate(rooms):
        x, y = room['center']
        w, h = room['width'], room['height']
        
        # Create box for each room
        fig.add_trace(go.Mesh3d(
            x=[x, x+w, x+w, x, x, x+w, x+w, x],
            y=[y, y, y+h, y+h, y, y, y+h, y+h],
            z=[0, 0, 0, 0, height, height, height, height],
            i=[0, 0, 0, 1, 4, 4, 4, 5],
            j=[1, 2, 4, 2, 5, 6, 7, 6],
            k=[2, 3, 7, 3, 6, 7, 6, 7],
            opacity=0.7,
            name=f'Room {i+1}'
        ))
    
    # Update layout
    fig.update_layout(
        scene=dict(
            aspectmode='data',
            camera=dict(
                up=dict(x=0, y=0, z=1),
                center=dict(x=0, y=0, z=0),
                eye=dict(x=1.5, y=1.5, z=1.5)
            )
        ),
        showlegend=True
    )
    
    return fig

# File uploader
uploaded_file = st.file_uploader("Upload 2D floor plan (PNG, JPG)", type=['png', 'jpg', 'jpeg'])

if uploaded_file is not None:
    # Convert uploaded file to 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)
    
    # Create and display 3D visualization
    st.subheader("3D Visualization")
    fig = create_3d_visualization(rooms)
    st.plotly_chart(fig, use_container_width=True)
    
    # Display room information
    st.subheader("Room Information")
    for i, room in enumerate(rooms):
        st.write(f"Room {i+1}:")
        st.write(f"- Area: {room['area']:.2f} square pixels")
        st.write(f"- Dimensions: {room['width']}x{room['height']} pixels")

st.sidebar.markdown("""
### Instructions
1. Upload a 2D floor plan image (PNG or JPG format)
2. The app will automatically:
   - Detect rooms and spaces
   - Create a 3D visualization
   - Show room information
3. Use mouse to rotate and zoom the 3D view
""")

st.sidebar.markdown("""
### Notes
- The floor plan should have clear walls and room boundaries
- White background with dark walls works best
- Simple floor plans produce better results
""")