Code Editor for app.py

import streamlit as st
import cv2
import numpy as np
from skimage.morphology import skeletonize
from skimage import measure
import plotly.graph_objects as go
import matplotlib.pyplot as plt
from utils import detect_walls, segment_rooms, create_3d_model

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

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

if uploaded_file is not None:
    # Convert uploaded file to OpenCV format
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    floor_plan = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    
    # Display original floor plan
    st.subheader("Original Floor Plan")
    st.image(floor_plan, channels="BGR")
    
    # Process floor plan
    with st.spinner("Processing floor plan..."):
        # Detect walls
        walls = detect_walls(floor_plan)
        st.subheader("Wall Detection")
        st.image(walls, caption="Detected Walls")
        
        # Segment rooms
        rooms, n_rooms = segment_rooms(walls)
        st.subheader("Room Segmentation")
        
        # Create color map for visualization
        colors = plt.cm.rainbow(np.linspace(0, 1, n_rooms))
        colored_rooms = np.zeros((*rooms.shape, 3))
        for i in range(n_rooms):
            colored_rooms[rooms == i] = colors[i][:3]
        st.image(colored_rooms, caption=f"Segmented Rooms (Found {n_rooms} rooms)")
        
        # Create 3D visualization
        fig = create_3d_model(walls, rooms)
        st.subheader("3D Visualization")
        st.plotly_chart(fig)
        
    st.success("Processing complete!")
    
st.markdown("""
### How it works:
1. **Wall Detection**: Uses edge detection and morphological operations to identify walls
2. **Room Segmentation**: Implements flood fill algorithm to identify separate rooms
3. **3D Visualization**: Extrudes the 2D plan to create a 3D model with proper wall heights
""")