trttung1610 commited on
Commit
1dbec60
·
1 Parent(s): 5a6b07a

Delete app_v1.py

Browse files
Files changed (1) hide show
  1. app_v1.py +0 -146
app_v1.py DELETED
@@ -1,146 +0,0 @@
1
- import gradio as gr
2
- from PIL import Image
3
- import torch
4
- import torchvision.models as models
5
- import torchvision.transforms as transforms
6
- import cv2
7
- import numpy as np
8
- import openpyxl
9
- import os
10
- from tkinter import filedialog
11
-
12
- # Load the pre-trained EfficientNet-B7 model
13
- model = models.efficientnet_b7(pretrained=True)
14
- model.eval()
15
-
16
- # Define the transformations to be applied to the input image
17
- transform = transforms.Compose([
18
- transforms.Resize((224, 224)),
19
- transforms.ToTensor(),
20
- transforms.Normalize(mean=[0.485, 0.456, 0.406],
21
- std=[0.229, 0.224, 0.225])
22
- ])
23
- def predict_house_area(excel_file, image_files):
24
- total_area_sqm = 0
25
- predicted_areas = []
26
-
27
- # Check if the excel_file is provided
28
- if excel_file is not None:
29
- # Load the existing Excel workbook
30
- workbook = openpyxl.load_workbook(excel_file.name)
31
- worksheet = workbook.active
32
- else:
33
- # Create a new Excel workbook
34
- workbook = openpyxl.Workbook()
35
- worksheet = workbook.active
36
-
37
- # Write the headers to the worksheet
38
- worksheet.cell(row=1, column=1).value = "Image File"
39
- worksheet.cell(row=1, column=2).value = "Predicted Area (sqm)"
40
-
41
- # Get the last row index to append new data
42
- last_row_index = worksheet.max_row if worksheet.max_row else 1
43
-
44
- # Loop over all the images
45
- for i, image_file in enumerate(image_files):
46
- # Load the input image
47
- img = Image.open(image_file.name)
48
- # Extract the image file name from the path
49
- image_file_name = os.path.basename(image_file.name)
50
- # Check if the image is PNG and convert to JPEG if it is
51
- if img.format == "PNG":
52
- # Convert the image to RGB format
53
- img = img.convert("RGB")
54
-
55
- # Apply the transformations to the input image
56
- img_transformed = transform(img)
57
-
58
- # Add a batch dimension to the transformed image tensor
59
- img_transformed_batch = torch.unsqueeze(img_transformed, 0)
60
-
61
- # Use the pre-trained model to make a prediction on the input image
62
- with torch.no_grad():
63
- output = model(img_transformed_batch)
64
-
65
- # Convert the output tensor to a probability distribution using softmax
66
- softmax = torch.nn.Softmax(dim=1)
67
- output_probs = softmax(output)
68
-
69
- # Extract the predicted class (house square footage) from the output probabilities
70
- predicted_class = torch.argmax(output_probs)
71
-
72
- # Calculate the predicted area based on the predicted class
73
- predicted_area_sqm = 0
74
- if predicted_class in [861, 648, 594, 894, 799, 896, 454]:
75
- # Convert to grayscale and apply adaptive thresholding
76
- gray = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY)
77
- gray = cv2.GaussianBlur(gray, (5, 5), 0)
78
- mask = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 11, 2)
79
-
80
- # Apply Canny edge detection to the binary mask
81
- edges = cv2.Canny(mask, 30, 100)
82
-
83
- # Apply dilation to fill gaps in the contour
84
- kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
85
- dilated = cv2.dilate(edges, kernel, iterations=2)
86
- eroded = cv2.erode(dilated, kernel, iterations=1)
87
-
88
- # Find contours in binary mask
89
- contours, _ = cv2.findContours(eroded, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
90
-
91
- # Find largest contour and calculate area
92
- max_area = 0
93
- for c in contours:
94
- area = cv2.contourArea(c)
95
- if area > max_area:
96
- max_area = area
97
-
98
- # Convert pixel area to square meters
99
- pixels_per_meter = 300 # adjust this value based on your image resolution and actual room dimensions
100
- predicted_area_sqm = (max_area + 10) / (2 * pixels_per_meter ** 2)
101
- else:
102
- predicted_area_sqft = predicted_class.item()
103
- predicted_area_sqm = predicted_area_sqft * 0.092903 / 4.2
104
-
105
- # Add the predicted area to the sum
106
- total_area_sqm += predicted_area_sqm
107
-
108
- # Add the predicted area to the list of predicted areas
109
- predicted_areas.append(predicted_area_sqm)
110
-
111
- # Write the room ID, image file name, and predicted area to the worksheet
112
- worksheet.cell(row=last_row_index + i + 1, column=1).value = image_file_name
113
- worksheet.cell(row=last_row_index + i + 1, column=2).value = predicted_area_sqm
114
-
115
- # Save the workbook to a temporary file
116
- temp_file = "predicted_areas.xlsx"
117
- workbook.save(temp_file)
118
-
119
- # Get the path of the first uploaded image
120
- first_image_path = image_files[0].name if image_files else None
121
-
122
- return f"Sum of predicted house square footage: {total_area_sqm:.2f} square meters", temp_file ,first_image_path
123
-
124
-
125
- inputs = [
126
- gr.inputs.File(label="Excel File", type="file"),
127
- gr.inputs.File(label="Images", type="file", file_count="multiple")
128
- ]
129
-
130
- outputs = [
131
- gr.outputs.Textbox(label="Sum of Predicted House Square Footage"),
132
- gr.outputs.File(label="Excel Result"),
133
- gr.outputs.Image(type="pil", label="Uploaded Image")
134
- ]
135
-
136
- interface = gr.Interface(
137
- fn=predict_house_area,
138
- inputs=inputs,
139
- outputs=outputs,
140
- title="House Predictor",
141
- allow_flagging="never" # Disable flag button
142
- )
143
-
144
- if __name__ == "__main__":
145
- interface.launch()
146
-