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

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)