import easyocr
from PIL import Image
import cv2
import re
import gradio as gr

# Initialize the OCR reader
reader = easyocr.Reader(['en'])

# Define patterns for different room types and dimensions
room_patterns = {
    'bedroom': r'bedroom|bed\s?rm',
    'bathroom': r'bathroom|bath\s?rm',
    'kitchen': r'kitchen',
    'living room': r'living\s?room|sitting\s?room',
    'dining room': r'dining\s?room',
    # Add more patterns as needed
}

dimension_pattern = r"(\d+'\s?\d+\")|(\d+\.?\d*\s?[x×]\s?\d+\.?\d*)"  # Pattern to match dimensions like 10'6" or 10x12

# Function to extract room data with counts and measurements
def extract_room_data(results, room_patterns, dimension_pattern):
    # ... (same as before) ...

# Function to process the uploaded image
def process_image(image):
    # Convert the Gradio image to PIL Image
    image = Image.fromarray(image)

    # Perform OCR
    results = reader.readtext(np.array(image), detail=0)

    # Extract room information
    room_data_with_counts = extract_room_data(results, room_patterns, dimension_pattern)

    # Format output for display
    output_text = "Extracted Room Data with Counts and Measurements:\n"
    for room_type, data in room_data_with_counts.items():
        output_text += f"- {room_type.capitalize()}:\n"
        output_text += f"  Count: {data['count']}\n"
        output_text += f"  Measurements: {data['measurements']}\n"

    return output_text

# Create the Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="pil"),  # Use type="pil" for PIL Image input
    outputs="text",
    title="Floor Plan Room Detection",
    description="Upload a floor plan image to extract room information."
)

# Launch the interface
iface.launch(share=True)  # Set share=True to get a shareable link