Image classification

Image classification assigns a label or class to an image. Unlike text or audio classification, the inputs are the pixel values that comprise an image. There are many applications for image classification such as detecting damage after a natural disaster, monitoring crop health, or helping screen medical images for signs of disease.

This guide will show you how to:

1. Finetune ViT on the Food-101 dataset to classify a food item in an image.
2. Use your finetuned model for inference.

Before you begin, make sure you have all the necessary libraries installed:

pip install transformers datasets evaluate

We encourage you to login to your Hugging Face account so you can upload and share your model with the community. When prompted, enter your token to login:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Load Food-101 dataset

Start by loading a smaller subset of the Food-101 dataset from the 🤗 Datasets library. This’ll give you a chance to experiment and make sure everythings works before spending more time training on the full dataset.

>>> from datasets import load_dataset

>>> food = load_dataset("food101", split="train[:5000]")

Split the dataset’s train split into a train and test set with the train_test_split method:

>>> food = food.train_test_split(test_size=0.2)

Then take a look at an example:

>>> food["train"][0]
{'image': <PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=512x512 at 0x7F52AFC8AC50>,
 'label': 79}

There are two fields:

• image: a PIL image of the food item.
• label: the label class of the food item.

To make it easier for the model to get the label name from the label id, create a dictionary that maps the label name to an integer and vice versa:

>>> labels = food["train"].features["label"].names
>>> label2id, id2label = dict(), dict()
>>> for i, label in enumerate(labels):
...     label2id[label] = str(i)
...     id2label[str(i)] = label

Now you can convert the label id to a label name:

>>> id2label[str(79)]
'prime_rib'

Preprocess

The next step is to load a ViT image processor to process the image into a tensor:

>>> from transformers import AutoImageProcessor

>>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")

Apply some image transformations to the images to make the model more robust against overfitting. Here you’ll use torchvision’s transforms module, but you can also use any image library you like.

Crop a random part of the image, resize it, and normalize it with the image mean and standard deviation:

>>> from torchvision.transforms import RandomResizedCrop, Compose, Normalize, ToTensor

>>> normalize = Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
>>> size = (
...     image_processor.size["shortest_edge"]
...     if "shortest_edge" in image_processor.size
...     else (image_processor.size["height"], image_processor.size["width"])
... )
>>> _transforms = Compose([RandomResizedCrop(size), ToTensor(), normalize])

Then create a preprocessing function to apply the transforms and return the pixel_values - the inputs to the model - of the image:

>>> def transforms(examples):
...     examples["pixel_values"] = [_transforms(img.convert("RGB")) for img in examples["image"]]
...     del examples["image"]
...     return examples

To apply the preprocessing function over the entire dataset, use 🤗 Datasets with_transform method. The transforms are applied on the fly when you load an element of the dataset:

>>> food = food.with_transform(transforms)

Now create a batch of examples using DataCollatorWithPadding. Unlike other data collators in 🤗 Transformers, the DefaultDataCollator does not apply additional preprocessing such as padding.

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator()

Evaluate

Including a metric during training is often helpful for evaluating your model’s performance. You can quickly load a evaluation method with the 🤗 Evaluate library. For this task, load the accuracy metric (see the 🤗 Evaluate quick tour to learn more about how to load and compute a metric):

>>> import evaluate

>>> accuracy = evaluate.load("accuracy")

Then create a function that passes your predictions and labels to compute to calculate the accuracy:

>>> import numpy as np


>>> def compute_metrics(eval_pred):
...     predictions = np.argmax(eval_pred.predictions, axis=1)
...     return accuracy.compute(predictions=predictions, references=eval_pred.label_ids)

Your compute_metrics function is ready to go now, and you’ll return to it when you setup your training.

Train

>>> from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

>>> model = AutoModelForImageClassification.from_pretrained(
...     "google/vit-base-patch16-224-in21k",
...     num_labels=len(labels),
...     id2label=id2label,
...     label2id=label2id,
... )

>>> training_args = TrainingArguments(
...     output_dir="my_awesome_food_model",
...     remove_unused_columns=False,
...     evaluation_strategy="epoch",
...     save_strategy="epoch",
...     learning_rate=5e-5,
...     per_device_train_batch_size=16,
...     gradient_accumulation_steps=4,
...     per_device_eval_batch_size=16,
...     num_train_epochs=3,
...     warmup_ratio=0.1,
...     logging_steps=10,
...     load_best_model_at_end=True,
...     metric_for_best_model="accuracy",
...     push_to_hub=True,
... )

>>> trainer = Trainer(
...     model=model,
...     args=training_args,
...     data_collator=data_collator,
...     train_dataset=food["train"],
...     eval_dataset=food["test"],
...     tokenizer=image_processor,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

>>> trainer.push_to_hub()

Inference

Great, now that you’ve finetuned a model, you can use it for inference!

Load an image you’d like to run inference on:

>>> ds = load_dataset("food101", split="validation[:10]")
>>> image = ds["image"][0]

The simplest way to try out your finetuned model for inference is to use it in a pipeline(). Instantiate a pipeline for image classification with your model, and pass your image to it:

>>> from transformers import pipeline

>>> classifier = pipeline("image-classification", model="my_awesome_food_model")
>>> classifier(image)
[{'score': 0.35574808716773987, 'label': 'beignets'},
 {'score': 0.018057454377412796, 'label': 'chicken_wings'},
 {'score': 0.017733804881572723, 'label': 'prime_rib'},
 {'score': 0.016335085034370422, 'label': 'bruschetta'},
 {'score': 0.0160061065107584, 'label': 'ramen'}]

You can also manually replicate the results of the pipeline if you’d like:

>>> from transformers import AutoImageProcessor
>>> import torch

>>> image_processor = AutoImageProcessor.from_pretrained("my_awesome_food_model")
>>> inputs = image_processor(image, return_tensors="pt")

>>> from transformers import AutoModelForImageClassification

>>> model = AutoModelForImageClassification.from_pretrained("my_awesome_food_model")
>>> with torch.no_grad():
...     logits = model(**inputs).logits

>>> predicted_label = logits.argmax(-1).item()
>>> model.config.id2label[predicted_label]
'beignets'