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()