timm_resnet50_image_classification_top5_predictions.py

1import torch
2import timm
3from PIL import Image
4from urllib.request import urlopen
5
6# 1. Create a model
7model = timm.create_model('resnet50', pretrained=True)
8model.eval()
9
10# 2. Get model specific transforms (normalization, resizing)
11data_config = timm.data.resolve_model_data_config(model)
12transforms = timm.data.create_transform(**data_config, is_training=False)
13
14# 3. Load and process an image
15img = Image.open(urlopen('https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/beignets-task-guide.png'))
16output = model(transforms(img).unsqueeze(0))  # unsqueeze to add batch dimension
17
18# 4. Get the top-5 predictions
19top5_probabilities, top5_class_indices = torch.topk(output.softmax(dim=1) * 100, k=5)
20
21print(top5_probabilities)
22print(top5_class_indices)