Views
No views yet
1import torch
2import torch.nn as nn
3from torchvision import models, transforms
4from PIL import Image
5
6# Load model
7model = models.convnext_base(weights=None)
8num_features = model.classifier[2].in_features
9model.classifier[2] = nn.Sequential(
10 nn.Dropout(0.6),
11 nn.Linear(num_features, 1024),
12 nn.ReLU(),
13 nn.Dropout(0.5),
14 nn.Linear(1024, 200)
15)
16
17# Load weights
18model.load_state_dict(torch.load('final_model.pth', map_location='cpu'))
19model.eval()
20
21# Preprocessing
22transform = transforms.Compose([
23 transforms.Resize((224, 224)),
24 transforms.ToTensor(),
25 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
26])
27
28# Predict
29image = Image.open('bird.jpg').convert('RGB')
30image_tensor = transform(image).unsqueeze(0)
31
32with torch.no_grad():
33 outputs = model(image_tensor)
34 probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
35 top5_prob, top5_indices = torch.topk(probabilities, 5)
36
37print("Top 5 Predictions:")
38for prob, idx in zip(top5_prob, top5_indices):
39 print(f"Class {idx}: {prob.item()*100:.2f}%")final_model.pth (1.06 GB): Full model weights@techreport{WahCUB_200_2011,
Title = {{The Caltech-UCSD Birds-200-2011 Dataset}},
Author = {Wah, C. and Branson, S. and Welinder, P. and Perona, P. and Belongie, S.},
Year = {2011},
Institution = {California Institute of Technology},
Number = {CNS-TR-2011-001}
}