| --- |
| license: mit |
| tags: |
| - floorplan |
| - real-estate |
| - image-classification |
| datasets: |
| - custom |
| --- |
| |
| # FloorplanValidator |
|
|
| This model distinguishes between floorplan images and non-floorplan images in real estate listings. |
|
|
| ## Model Details |
|
|
| - Model type: ResNet50 fine-tuned for binary classification |
| - Task: Binary image classification |
| - Training data: Custom dataset of floorplan and non-floorplan images |
| - Class labels: 0 (floorplan), 1 (no_image) |
| |
| ## Intended Use |
| |
| - Identify valid floorplan images in real estate listings |
| - Filter out non-floorplan images |
| |
| ## Usage |
| |
| ```python |
| import torch |
| import torch.nn as nn |
| from torchvision import transforms, models |
| from huggingface_hub import hf_hub_download |
| from PIL import Image |
|
|
| # 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 |
| transform = transforms.Compose([ |
| transforms.Resize((224, 224)), |
| transforms.ToTensor(), |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) |
| ]) |
| |
| # Make a prediction |
| image = Image.open("your_image.jpg").convert('RGB') |
| input_tensor = transform(image).unsqueeze(0) |
| |
| with torch.no_grad(): |
| output = model(input_tensor) |
| probs = torch.softmax(output, dim=1) |
| pred_class = torch.argmax(probs, dim=1).item() |
| confidence = probs[0][pred_class].item() |
| |
| result = { |
| 'class': "floorplan" if pred_class == 0 else "non-floorplan", |
| 'confidence': confidence |
| } |
| print(result) |
| ``` |
| |