Code Editor for app.py

import streamlit as st
import numpy as np
import cv2
import matplotlib.pyplot as plt
from skimage.segmentation import watershed
from skimage.feature import peak_local_max
from scipy import ndimage
import plotly.graph_objects as go
from pathlib import Path
import io
from PIL import Image

st.title("2D Floor Plan to 3D Converter")
st.write("Upload a 2D floor plan image to convert it into a 3D visualization")

def detect_walls(image):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Apply adaptive thresholding
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                 cv2.THRESH_BINARY_INV, 11, 2)
    
    # Find contours
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Draw walls
    wall_mask = np.zeros_like(gray)
    cv2.drawContours(wall_mask, contours, -1, (255), 2)
    
    return wall_mask

def segment_rooms(image, wall_mask):
    # Distance transform
    dist_transform = cv2.distanceTransform(255 - wall_mask, cv2.DIST_L2, 5)
    
    # Find local maxima
    local_max = peak_local_max(dist_transform, indices=False, min_distance=20,
                              labels=255 - wall_mask)
    
    # Create markers for watershed
    markers = ndimage.label(local_max)[0]
    
    # Apply watershed
    labels = watershed(-dist_transform, markers, mask=255 - wall_mask)
    
    return labels

def create_3d_visualization(labels, wall_mask):
    height = 3.0  # Standard room height
    
    # Create coordinates for walls and rooms
    y, x = np.nonzero(wall_mask)
    z_walls = np.ones_like(y) * height
    
    # Create room floors
    rooms_y, rooms_x = np.nonzero(labels > 0)
    z_rooms = np.zeros_like(rooms_y)
    
    # Combine coordinates
    x_coords = np.concatenate([x, rooms_x])
    y_coords = np.concatenate([y, rooms_y])
    z_coords = np.concatenate([z_walls, z_rooms])
    
    # Create color array
    colors = np.zeros_like(z_coords)
    colors[:len(z_walls)] = 1  # Walls in one color
    colors[len(z_walls):] = 2  # Floors in another color
    
    # Create 3D scatter plot
    fig = go.Figure(data=[go.Scatter3d(
        x=x_coords,
        y=y_coords,
        z=z_coords,
        mode='markers',
        marker=dict(
            size=2,
            color=colors,
            colorscale=[[0, 'rgb(200,200,200)'],
                       [0.5, 'rgb(100,100,100)'],
                       [1, 'rgb(50,50,50)']],
            opacity=0.8
        )
    )])
    
    # Update layout
    fig.update_layout(
        scene=dict(
            aspectmode='data',
            xaxis_title='X',
            yaxis_title='Y',
            zaxis_title='Height'
        ),
        margin=dict(l=0, r=0, b=0, t=0)
    )
    
    return fig

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

if uploaded_file is not None:
    # Read and process image
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    
    # Create columns for displaying results
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Original Floor Plan")
        st.image(image, channels="BGR", use_column_width=True)
    
    # Detect walls
    wall_mask = detect_walls(image)
    
    with col2:
        st.subheader("Detected Walls")
        st.image(wall_mask, use_column_width=True)
    
    # Segment rooms
    room_labels = segment_rooms(image, wall_mask)
    
    st.subheader("Room Segmentation")
    fig_rooms, ax = plt.subplots()
    ax.imshow(room_labels, cmap='nipy_spectral')
    st.pyplot(fig_rooms)
    
    # Create 3D visualization
    st.subheader("3D Visualization")
    fig_3d = create_3d_visualization(room_labels, wall_mask)
    st.plotly_chart(fig_3d, use_container_width=True)
    
    st.info("""
    **How it works:**
    1. Wall Detection: Uses adaptive thresholding and contour detection
    2. Room Segmentation: Applies watershed algorithm to segment different rooms
    3. 3D Visualization: Creates a 3D model using detected walls and room segments
    
    Note: The quality of the conversion depends on the clarity and simplicity of the input floor plan.
    """)