Code Editor for app.py

import streamlit as st
import cv2
import numpy as np
from PIL import Image
import io

def convert_to_cad_style(image):
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    
    # Apply Gaussian blur to reduce noise
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    
    # Edge detection using Canny
    edges = cv2.Canny(blurred, 50, 150)
    
    # Dilate the edges to make them more prominent
    kernel = np.ones((3,3), np.uint8)
    dilated = cv2.dilate(edges, kernel, iterations=1)
    
    # Create a white background
    result = np.ones_like(image) * 255
    
    # Draw the edges in black on the white background
    result[dilated > 0] = [0, 0, 0]
    
    return result

st.title("Sketch to CAD-Style Floor Plan Converter")

st.write("""
Upload a hand-drawn floor plan sketch and this app will convert it into a CAD-style drawing.
For best results:
- Use clear, dark lines in your sketch
- Ensure good contrast between lines and background
- Avoid unnecessary details
""")

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

if uploaded_file is not None:
    # Convert the uploaded file to an OpenCV image
    file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    
    # Create columns for before/after comparison
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Original Sketch")
        st.image(image, channels="BGR", use_column_width=True)
    
    # Convert the sketch to CAD style
    cad_style = convert_to_cad_style(image)
    
    with col2:
        st.subheader("CAD-Style Output")
        st.image(cad_style, channels="BGR", use_column_width=True)
    
    # Add download button for the processed image
    processed_img = Image.fromarray(cv2.cvtColor(cad_style, cv2.COLOR_BGR2RGB))
    buf = io.BytesIO()
    processed_img.save(buf, format="PNG")
    st.download_button(
        label="Download CAD-Style Image",
        data=buf.getvalue(),
        file_name="cad_style_floor_plan.png",
        mime="image/png"
    )