Square_Predict / app.py
TuanScientist's picture
Update app.py
156e2ab verified
import gradio as gr
from PIL import Image
import torch
import torchvision
from torchvision.models import EfficientNet_B7_Weights
import torchvision.models as models
import torchvision.transforms as transforms
import cv2
import numpy as np
import openpyxl
import os
# Load the pre-trained EfficientNet-B7 model
model = torchvision.models.efficientnet_b7(weights=EfficientNet_B7_Weights.DEFAULT)
model.eval()
# Define the transformations to be applied to the input image
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
def predict_house_area(image_file):
total_area_sqm = 0
predicted_areas = []
# Load the input image from the Gradio Image component
img = Image.fromarray(image_file)
image_file_name = "single_image.jpg"
if img.format == "PNG":
img = img.convert("RGB")
img_transformed = transform(img)
img_transformed_batch = torch.unsqueeze(img_transformed, 0)
with torch.no_grad():
output = model(img_transformed_batch)
softmax = torch.nn.Softmax(dim=1)
output_probs = softmax(output)
predicted_class = torch.argmax(output_probs)
predicted_area_sqm = 0
if predicted_class in [861, 648, 594, 894, 799, 896, 454]:
gray = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY)
gray = cv2.GaussianBlur(gray, (5, 5), 0)
mask = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 11, 2)
edges = cv2.Canny(mask, 30, 100)
kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
dilated = cv2.dilate(edges, kernel, iterations=2)
eroded = cv2.erode(dilated, kernel, iterations=1)
contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
max_area = 0
for c in contours:
area = cv2.contourArea(c)
if area > max_area:
max_area = area
pixels_per_meter = 300
predicted_area_sqm = (max_area + 10) / (2 * pixels_per_meter ** 2)
else:
predicted_area_sqft = predicted_class.item()
predicted_area_sqm = predicted_area_sqft * 0.092903 / 4.2
total_area_sqm += predicted_area_sqm
predicted_areas.append(predicted_area_sqm)
workbook = openpyxl.Workbook()
worksheet = workbook.active
worksheet.cell(row=1, column=1).value = "Image File"
worksheet.cell(row=1, column=2).value = "Predicted Area (sqm)"
worksheet.cell(row=2, column=1).value = image_file_name
worksheet.cell(row=2, column=2).value = predicted_area_sqm
temp_file = "predicted_area.xlsx"
workbook.save(temp_file)
return f"Predicted house square footage: {predicted_area_sqm:.2f} square meters", temp_file
inputs = [
gr.Image(label="Image")
]
outputs = [
gr.Textbox(label="Predicted Room Square"),
gr.File(label="Excel Printed Result"),
]
interface = gr.Interface(
fn=predict_house_area,
inputs=inputs,
outputs=outputs,
title="Room Square Predictor",
allow_flagging="never" # Disable flag button
)
if __name__ == "__main__":
interface.launch()