Id2Labels - how to map classification result to labels #1982
Replies: 2 comments 3 replies
-
You can map the class index to the class name by downloading the imagenet-1k or imagenet-21k class-name list depending on the dataset which model are finetuned on. Here is a code snippet for imagenet-1k (see this): import torch
import timm
import urllib
from PIL import Image
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform
# step 1. create model and transform.
model = timm.create_model('efficientnet_b0', pretrained=True)
model.eval()
config = resolve_data_config({}, model=model)
transform = create_transform(**config)
# step 2. download image and convert it to tensor.
url, filename = ("https://github.com/pytorch/hub/raw/master/images/dog.jpg", "dog.jpg")
urllib.request.urlretrieve(url, filename)
img = Image.open(filename).convert('RGB')
tensor = transform(img).unsqueeze(0) # transform and add batch dimension
# step 3. inference model.
with torch.no_grad():
out = model(tensor)
probabilities = torch.nn.functional.softmax(out[0], dim=0)
# step 4. download imagenet-1k label.
url, filename = ("https://raw.githubusercontent.com/pytorch/hub/master/imagenet_classes.txt", "imagenet_classes.txt")
urllib.request.urlretrieve(url, filename)
with open("imagenet_classes.txt", "r") as f:
categories = [s.strip() for s in f.readlines()]
# step 5. print class names and probabilities like:
# [('Samoyed', 0.6425196528434753), ('Pomeranian', 0.04062102362513542), ('keeshond', 0.03186424449086189), ('white wolf', 0.01739676296710968), ('Eskimo dog', 0.011717947199940681)]
top5_prob, top5_catid = torch.topk(probabilities, 5)
for i in range(top5_prob.size(0)):
print(categories[top5_catid[i]], top5_prob[i].item()) +) You can find the dataset name the model is trained on by checking the I hope this could help you. Thank you. Hankyul. |
Beta Was this translation helpful? Give feedback.
-
@cedricpinson @hankyul2 The best approach for timm models is what I put together in the API inferencing code (for the HF hub timm classification widgets). It uses metadata builtin to this repo about different classifier layouts and has detailed class description support (beyond most common impl). The |
Beta Was this translation helpful? Give feedback.
-
Hello,
I am pretty new and my question will probably be stupid. I am trying different models to classify images.
I am using different VIT models with this code:
When testing some examples on hugging face I was able to try some models that contains those labels, but not always
Is there something I missing or a way to fetch id2labels for all those models ?
Beta Was this translation helpful? Give feedback.
All reactions