Views
No views yet
1import torch
2import torchvision.models as models
3import torch.nn as nn
4from PIL import Image
5from torchvision import transforms
6from huggingface_hub import hf_hub_download
7
8# Load model
9model = models.resnet18(pretrained=False)
10num_features = model.fc.in_features
11model.fc = nn.Sequential(
12 nn.Linear(num_features, 256),
13 nn.ReLU(),
14 nn.Dropout(0.3),
15 nn.Linear(256, 10)
16)
17
18# Download and load weights
19model_path = hf_hub_download(repo_id="sabarsbb/cifar10-image-classifier-v1", filename="best_model.pth")
20checkpoint = torch.load(model_path, map_location='cpu')
21model.load_state_dict(checkpoint['model_state_dict'])
22model.eval()
23
24# Inference
25transform = transforms.Compose([
26 transforms.Resize((224, 224)),
27 transforms.ToTensor(),
28 transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
29])
30
31image = Image.open('your_image.jpg')
32image_tensor = transform(image).unsqueeze(0)
33
34with torch.no_grad():
35 outputs = model(image_tensor)
36 _, predicted = torch.max(outputs, 1)
37
38classes = ['airplane', 'automobile', 'bird', 'cat', 'deer',
39 'dog', 'frog', 'horse', 'ship', 'truck']
40print(f"Predicted class: {classes[predicted.item()]}")