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