timm_resnet50_image_classification_top5_predictions.py

1import timm
2import torch
3from PIL import Image
4from timm.data import resolve_data_config
5from timm.data.transforms_factory import create_transform
6
7# 1. Create the model
8model = timm.create_model('resnet50', pretrained=True)
9model.eval()
10
11# 2. Configure the data transforms
12config = resolve_data_config({}, model=model)
13transform = create_transform(**config)
14
15# 3. Load and preprocess an image (using a placeholder or local path)
16# Assuming 'dog.jpg' exists in the current directory
17img = Image.open('dog.jpg').convert('RGB')
18tensor = transform(img).unsqueeze(0) # transform and add batch dimension
19
20# 4. Perform inference
21with torch.no_grad():
22    out = model(tensor)
23
24# 5. Get top-5 probabilities and class indices
25probabilities = torch.nn.functional.softmax(out[0], dim=0)
26top5_prob, top5_catid = torch.topk(probabilities, 5)
27
28for i in range(top5_prob.size(0)):
29    print(f"Class ID: {top5_catid[i].item()}, Probability: {top5_prob[i].item():.4f}")