File size: 2,120 Bytes
dae751f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ad7a20d
 
 
dae751f
 
ad7a20d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dae751f
 
ad7a20d
dae751f
 
 
 
 
 
ad7a20d
dae751f
 
 
 
 
ad7a20d
 
 
dae751f
 
ad7a20d
dae751f
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
---
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)
```