taarhissian's picture
Update app.py
e5e2128 verified
raw
history blame
1.8 kB
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