init commit

app.py

@@ -0,0 +1,93 @@
+import torch
+import torch.nn as nn
+import torchvision.models as models
+import torchvision.transforms as transforms
+import torch.nn.functional as F
+from PIL import Image
+import gradio as gr
+
+# Define your model architecture
+class EfficientNetMultiTask(nn.Module):
+    def __init__(self, n_area_classes, n_room_classes):
+        super(EfficientNetMultiTask, self).__init__()
+        self.efficientnet = models.efficientnet_b0(pretrained=False)
+        in_features = self.efficientnet.classifier[1].in_features
+        self.area_classifier = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, n_area_classes)
+        )
+        self.room_classifier = nn.Sequential(
+            nn.Linear(in_features, 512),
+            nn.ReLU(),
+            nn.Dropout(0.3),
+            nn.Linear(512, n_room_classes)
+        )
+        self.efficientnet.classifier = nn.Identity()
+
+    def forward(self, x):
+        features = self.efficientnet(x)
+        area_pred = self.area_classifier(features)
+        room_pred = self.room_classifier(features)
+        return area_pred, room_pred
+
+
+# Load model
+n_area_classes = 21  # Adjust according to your area bins
+n_room_classes = 16  # Adjust based on your dataset
+model = EfficientNetMultiTask(n_area_classes=n_area_classes, n_room_classes=n_room_classes)
+
+# Load weights (ensure floorplan_model_classification.pth is in the same directory as app.py)
+model_weights_path = 'floorplan_model_classification.pth'  # Adjust with your model weights path
+model.load_state_dict(torch.load(model_weights_path, map_location=torch.device('cpu')))
+model.eval()
+
+# Define transformations
+test_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])
+])
+
+# Define area bins
+area_bins = [i for i in range(0, 525, 25)]  # [0, 25, ..., 500]
+area_bins.append(float('inf'))  # Add infinity for 500+ category
+
+def get_area_from_bin(area_bin_idx):
+    if area_bin_idx < len(area_bins) - 2:
+        return f"{area_bins[area_bin_idx]} - {area_bins[area_bin_idx + 1]} m²"
+    else:
+        return f"{area_bins[-2]}+ m²"
+
+# Prediction function
+def predict(image):
+    image = Image.fromarray(image).convert('RGB')
+    image = test_transform(image).unsqueeze(0)
+
+    with torch.no_grad():
+        area_output, room_output = model(image)
+        area_probabilities = F.softmax(area_output, dim=1)
+        room_probabilities = F.softmax(room_output, dim=1)
+        area_pred_idx = torch.argmax(area_probabilities, dim=1).item()
+        room_pred_idx = torch.argmax(room_probabilities, dim=1).item()
+        predicted_area = get_area_from_bin(area_pred_idx)
+        predicted_rooms = room_pred_idx + 1  # Adjusting back to original room labels
+
+    return predicted_area, str(predicted_rooms)
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="numpy", label="Upload Floor Plan Image"),  # Correct input
+    outputs=[
+        gr.Textbox(label="Predicted Total Area"),  # Correct output
+        gr.Textbox(label="Predicted Number of Rooms")
+    ],
+    title="Floor Plan Area and Room Predictor",
+    description="Upload a floor plan image, and the model will predict the total area range and the number of rooms."
+)
+
+# Launch Gradio interface
+if __name__ == "__main__":
+    interface.launch()

floorplan_model_classification.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a37c5986f0e9b486c8b2db8032630c581ee1b444e8cd80f7ace1f83aef6f1bb
+size 21671252

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+pillow
+gradio