Id2Labels - how to map classification result to labels

[cedricpinson]

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:

from pprint import pprint
from PIL import Image
import requests
import timm
import torch
import time
import sys

def execute_timm(pretrained, image, num_results=10):
    global label21k, label1k

    perf = time.perf_counter_ns()

    model = timm.create_model(pretrained, pretrained=True)

    if "id2label" not in model.pretrained_cfg:
        print("no lables found")
        pprint(model.pretrained_cfg)
        return {}

    # model = timm.create_model("hf_hub:timm/vit_large_patch14_clip_224", pretrained=True)
    # model = timm.create_model("hf_hub:timm/vit_large_patch14_clip_224.openai_ft_in12k_in1k", pretrained=True)
    model = model.eval()

    # get model specific transforms (normalization, resize)
    data_config = timm.data.resolve_data_config(model.pretrained_cfg, model=model)
    transform = timm.data.transforms_factory.create_transform(**data_config)

    output = model(transform(image).unsqueeze(0))  # unsqueeze single image into batch of 1

    print(model.pretrained_cfg.keys())
    # labels = model.pretrained_cfg['labels']
    top_k = num_results

    # Apply softmax to get predicted probabilities for each class
    probabilities = torch.nn.functional.softmax(output[0], dim=0)

    # Grab the values and indices of top 5 predicted classes
    values, indices = torch.topk(probabilities, top_k)

    id2label = model.pretrained_cfg["labels"]

    print(values, indices)
    # Prepare a nice dict of top k predictions
    results = [
        # {"label": labels[i], "score": v.item()}
        {"label": id2label[i], "value": v.item()}
        for i, v in zip(indices, values)
    ]

    perf = (time.perf_counter_ns() - perf) / 1e6
    return {"time": perf, "model": pretrained, "results": results}

if __name__ == "__main__":

    model = "vit_huge_patch14_clip_224.laion2b_ft_in12k"
    image = None

    if "--list" in sys.argv:
        pprint(timm.list_models(pretrained=True))
    else:
        if  "--model" in sys.argv:
            model = sys.argv[sys.argv.index("--model") + 1]

        if "--image" in sys.argv:
            image = sys.argv[sys.argv.index("--image") + 1]

            image = Image.open( "./test.jpg" )
            pprint(execute_timm(model, image))

When testing some examples on hugging face I was able to try some models that contains those labels, but not always

    processor = AutoImageProcessor.from_pretrained(pretrained)
    model = AutoModelForImageClassification.from_pretrained(pretrained)

    #inputs = processor(images=images, return_tensors="pt")
    inputs = processor(images[0], return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    # print(logits)
    top10 = logits[0].topk(num_results)

    results = []
    for i in range(num_results):
        class_idx = top10.indices[i].item()
        results.append(
            {"label": model.config.id2label[class_idx], "value": top10.values[i].item()}
        )

Is there something I missing or a way to fetch id2labels for all those models ?

[hankyul2]

Hi @cedricpinson

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 model_name.pretrained_tag. For example, in vit_xxx.openai_ft_in12k_in1k, ft_in1k means the model is finetuned on the imagenet-1k dataset.

I hope this could help you. Thank you.

Hankyul.

[cedricpinson]

Thank you 🙏

[rwightman]

@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).

https://github.com/huggingface/api-inference-community/blob/main/docker_images/timm/app/pipelines/image_classification.py

It uses metadata builtin to this repo about different classifier layouts and has detailed class description support (beyond most common impl). The inference.py script in this repo also does something similar...

[hankyul2]

Hi @rwightman

Thank you for your quick reply.

I just want to share the short code snippet about how to use the index_to_description API for converting model prediction to the label description. (refer to the inference.py in this repo)

import timm
import torch

model = timm.create_model('resnet50')
imagenet_subset = timm.data.infer_imagenet_subset(model)
dataset_info = timm.data.ImageNetInfo(imagenet_subset)
to_label = lambda x: dataset_info.index_to_description(x)
x = torch.rand(1, 3, 224, 224)
y = model(x).topk(5)[1][0]
print([to_label(i) for i in y])
# >> ['mountain bike, all-terrain bike, off-roader', 'wok', 'bullfrog, Rana catesbeiana', 'typewriter keyboard', 
#        'coho, cohoe, coho salmon, blue jack, silver salmon, Oncorhynchus kisutch']

Thank you.

Hankyul

[cedricpinson]

Thank you all for the answer 🙏