Classificação de imagens

A classificação de imagens atribui um rótulo ou classe a uma imagem. Diferentemente da classificação de texto ou áudio, as entradas são os valores de pixels que compõem uma imagem. Existem muitas aplicações para a classificação de imagens, como detectar danos após um desastre natural, monitorar a saúde das culturas ou ajudar a triar imagens médicas em busca de sinais de doenças.

Este guia ilustra como:

1. Ajuste fino de ViT no conjunto de dados Food-101 para classificar um item de comida em uma imagem.
2. Use seu modelo ajustado para inferência.

Para ver todas as arquiteturas e checkpoints compatíveis com esta tarefa, recomendamos consultar a página da tarefa

Antes de começar, verifique se você tem todas as bibliotecas necessárias instaladas:

pip install transformers datasets evaluate accelerate pillow torchvision scikit-learn trackio

Recomendamos que você faça login na sua conta do Hugging Face para enviar e compartilhar seu modelo com a comunidade. Quando solicitado, insira seu token para fazer login:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Carregar o conjunto de dados Food-101

Comece carregando um subconjunto menor do conjunto de dados Food-101 a partir da biblioteca 🤗 Datasets. Isso lhe dará a oportunidade de experimentar e garantir que tudo funcione antes de gastar mais tempo treinando no conjunto de dados completo.

>>> from datasets import load_dataset

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

Divida o split train do conjunto de dados em um conjunto de treinamento e um conjunto de teste com o método train_test_split:

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

Em seguida, veja um exemplo:

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

Cada exemplo no conjunto de dados tem dois campos:

• image: uma imagem PIL do item de comida
• label: a classe de rótulo do item de comida

Para facilitar a obtenção do nome do rótulo a partir do id do rótulo pelo modelo, crie um dicionário que mapeia o nome do rótulo para um inteiro e 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

Agora você pode converter o id do rótulo em um nome de rótulo:

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

Pré-processar

O próximo passo é carregar um processador de imagens ViT para processar a imagem em um tensor:

>>> from transformers import AutoImageProcessor

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

Aplique algumas transformações de imagem às imagens para tornar o modelo mais robusto contra overfitting. Aqui, você usará o módulo transforms do torchvision, mas também pode usar qualquer biblioteca de imagem que preferir.

Recorte uma parte aleatória da imagem, redimensione-a e normalize-a com a média e o desvio padrão da imagem:

>>> 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])

Em seguida, crie uma função de pré-processamento para aplicar as transformações e retornar os pixel_values — as entradas do modelo — da imagem:

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

Para aplicar a função de pré-processamento em todo o conjunto de dados, use o método with_transform do 🤗 Datasets. As transformações são aplicadas em tempo de execução ao carregar um elemento do conjunto de dados:

>>> food = food.with_transform(transforms)

Agora, crie um lote de exemplos usando DefaultDataCollator. Diferentemente de outros data collators no 🤗 Transformers, o DefaultDataCollator não aplica pré-processamento adicional, como padding.

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator()

Avaliar

Incluir uma métrica durante o treinamento é frequentemente útil para avaliar o desempenho do seu modelo. Você pode carregar rapidamente um método de avaliação com a biblioteca 🤗 Evaluate. Para esta tarefa, carregue a métrica accuracy (veja o quick tour do 🤗 Evaluate para saber mais sobre como carregar e calcular uma métrica):

>>> import evaluate

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

Em seguida, crie uma função que passa suas previsões e rótulos para compute para calcular a precisão:

>>> import numpy as np


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

Sua função compute_metrics está pronta e você voltará a ela ao configurar o treinamento.

Treinar

Se você não está familiarizado com o ajuste fino de um modelo usando o Trainer, confira o tutorial básico aqui!

Agora você está pronto para começar o treinamento do seu modelo! Carregue o ViT com AutoModelForImageClassification. Especifique o número de rótulos, junto com o número de rótulos esperados e os mapeamentos de rótulos:

>>> from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

>>> model = AutoModelForImageClassification.from_pretrained(
...     checkpoint,
...     num_labels=len(labels),
...     id2label=id2label,
...     label2id=label2id,
... )

Neste ponto, restam apenas três etapas:

1. Defina os hiperparâmetros de treinamento em TrainingArguments. É importante não remover colunas não utilizadas, pois isso descartaria a coluna image. Sem a coluna image, não é possível criar pixel_values. Defina remove_unused_columns=False para evitar esse comportamento! O único outro parâmetro obrigatório é output_dir, que especifica onde salvar o modelo. Você enviará este modelo para o Hub definindo push_to_hub=True (é necessário estar conectado ao Hugging Face para enviar o modelo). No final de cada epoch, o Trainer avaliará a precisão e salvará o checkpoint de treinamento.
2. Passe os argumentos de treinamento para o Trainer junto com o modelo, o conjunto de dados, o tokenizer, o data collator e a função compute_metrics.
3. Chame train() para fazer o ajuste fino do seu modelo.

>>> training_args = TrainingArguments(
...     output_dir="my_awesome_food_model",
...     remove_unused_columns=False,
...     eval_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_steps=0.1,
...     logging_steps=10,
...     report_to="trackio",
...     run_name="food101",
...     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"],
...     processing_class=image_processor,
...     compute_metrics=compute_metrics,
... )

>>> trainer.train()

Quando o treinamento for concluído, compartilhe seu modelo no Hub com o método push_to_hub() para que todos possam usá-lo:

>>> trainer.push_to_hub()

Para um exemplo mais detalhado de como fazer o ajuste fino de um modelo para classificação de imagens, confira o notebook do PyTorch correspondente.

Inferência

Ótimo, agora que você ajustou um modelo, pode usá-lo para inferência!

Carregue uma imagem na qual você deseja executar inferência:

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

A maneira mais simples de testar seu modelo ajustado para inferência é usá-lo em um pipeline(). Instancie um pipeline para classificação de imagem com seu modelo e passe sua imagem para ele:

>>> from transformers import pipeline

>>> classifier = pipeline("image-classification", model="my_awesome_food_model")
>>> classifier(image)
[{'score': 0.31856709718704224, 'label': 'beignets'},
 {'score': 0.015232225880026817, 'label': 'bruschetta'},
 {'score': 0.01519392803311348, 'label': 'chicken_wings'},
 {'score': 0.013022331520915031, 'label': 'pork_chop'},
 {'score': 0.012728818692266941, 'label': 'prime_rib'}]

Você também pode replicar manualmente os resultados do pipeline, se desejar:

Carregue um processador de imagem para pré-processar a imagem e retornar o input como tensores do PyTorch:

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

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

Passe suas entradas para o modelo e retorne os logits:

>>> from transformers import AutoModelForImageClassification

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

Obtenha o rótulo previsto com a maior probabilidade e use o mapeamento id2label do modelo para convertê-lo em um rótulo:

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