Upload FloorplanValidator model

README.md

@@ -27,36 +27,53 @@ This model distinguishes between floorplan images and non-floorplan images in re
 ## Usage
 
 ```python
-from huggingface_hub import hf_hub_download
 import torch
+import torch.nn as nn
+from torchvision import transforms, models
+from huggingface_hub import hf_hub_download
 from PIL import Image
-from torchvision import transforms
 
-# Load model
-model_path = hf_hub_download("acd20000/FloorplanValidator", "model.pt")
-model = torch.load(model_path, map_location=torch.device('cpu'))
+# Define the model architecture
+class RealEstateClassifier(nn.Module):
+    def __init__(self):
+        super().__init__()
+        # Load ResNet50
+        self.model = models.resnet50(pretrained=False)
+        # Modify final layer for binary classification
+        num_ftrs = self.model.fc.in_features
+        self.model.fc = nn.Linear(num_ftrs, 2)  # 2 classes: floorplan and no_image
+    
+    def forward(self, x):
+        return self.model(x)
+
+# Load the state dict
+model_path = hf_hub_download("acd20000/FloorplanValidator", "best_floorplan_classifier.pt")
+state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+
+# Create model and load weights
+model = RealEstateClassifier()
+model.load_state_dict(state_dict)
 model.eval()
 
-# Define transformation for input images
+# Define transformation
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
 ])
 
-# Load and transform an image
+# Make a prediction
 image = Image.open("your_image.jpg").convert('RGB')
 input_tensor = transform(image).unsqueeze(0)
 
-# Make prediction
 with torch.no_grad():
     output = model(input_tensor)
-    probabilities = torch.softmax(output, dim=1)
-    predicted_class = torch.argmax(probabilities, dim=1).item()
-    confidence = probabilities[0][predicted_class].item()
+    probs = torch.softmax(output, dim=1)
+    pred_class = torch.argmax(probs, dim=1).item()
+    confidence = probs[0][pred_class].item()
     
     result = {
-        'class': "floorplan" if predicted_class == 0 else "non-floorplan",
+        'class': "floorplan" if pred_class == 0 else "non-floorplan",
         'confidence': confidence
     }
     print(result)

best_floorplan_classifier.pt

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