Spaces:
Sleeping
Sleeping
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 |