File size: 3,349 Bytes
84d6f96 63adb04 84d6f96 63adb04 84d6f96 63adb04 84d6f96 63adb04 84d6f96 63adb04 84d6f96 63adb04 84d6f96 63adb04 84d6f96 |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import requests
import torch
from torchvision import transforms, models
import torch.nn as nn
from PIL import Image
import io
def classify_image_huggingface(image_path, repo_id="mertincesu/property-room-classifier"):
# Automatically select the best available device
device = torch.device("cuda" if torch.cuda.is_available() else
"mps" if torch.backends.mps.is_available() else
"cpu")
print(f"Using device: {device}")
# Load model weights directly from Hugging Face
model_url = f"https://huggingface.co/{repo_id}/resolve/main/pytorch_model.bin"
try:
# Download the model file
print(f"Downloading model from {model_url}")
response = requests.get(model_url)
response.raise_for_status()
# Load the model weights
model_binary = io.BytesIO(response.content)
state_dict = torch.load(model_binary, map_location=device)
# Extract model state dict and label mapping
if 'model_state_dict' in state_dict:
model_state = state_dict['model_state_dict']
label_to_index = state_dict.get('label_to_index', {
'bath': 0,
'bed': 1,
'din': 2,
'kitchen': 3,
'living': 4
})
else:
model_state = state_dict
label_to_index = {
'bath': 0,
'bed': 1,
'din': 2,
'kitchen': 3,
'living': 4
}
# Determine number of classes from weights
num_classes = model_state['fc.weight'].shape[0]
# Define model structure dynamically
model = models.resnet50(weights=None)
num_ftrs = model.fc.in_features
model.fc = nn.Linear(num_ftrs, num_classes)
# Load the model weights
model.load_state_dict(model_state)
model.to(device)
model.eval()
# Define image 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 process the image
img = Image.open(image_path).convert('RGB')
img_tensor = transform(img).unsqueeze(0).to(device)
# Make prediction
with torch.no_grad():
outputs = model(img_tensor)
probs = torch.nn.functional.softmax(outputs, dim=1)
_, predicted = torch.max(outputs, 1)
# Reverse the label mapping
index_to_label = {v: k for k, v in label_to_index.items()}
result = {
'class': index_to_label[predicted.item()],
'confidence': probs[0][predicted.item()].item() * 100
}
print(f"Class: {result['class']}")
print(f"Confidence: {result['confidence']:.2f}%")
return result
except Exception as e:
print(f"Error: {e}")
return None
if __name__ == "__main__":
import sys
if len(sys.argv) > 1:
image_path = sys.argv[1]
else:
image_path = input("Enter path to image: ")
classify_image_huggingface(image_path) |