Clasificación de imágenes

La clasificación de imágenes asigna una etiqueta o clase a una imagen. A diferencia de la clasificación de texto o audio, las entradas son los valores de píxeles que componen una imagen. Existen muchas aplicaciones para la clasificación de imágenes, como detectar daños después de un desastre natural, monitorear la salud de los cultivos o ayudar a examinar imágenes médicas en busca de signos de enfermedad.

Esta guía ilustra cómo:

1. Ajusta ViT sobre el conjunto de datos Food-101 para clasificar un elemento de comida en una imagen.
2. Usa tu modelo afinado para la inferencia.

Para ver todas las arquitecturas y puntos de control compatibles con esta tarea, recomendamos consultar la página de la tarea

Antes de comenzar, asegúrate de tener instaladas todas las bibliotecas necesarias:

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

Te animamos a iniciar sesión en tu cuenta de Hugging Face para subir y compartir tu modelo con la comunidad. Cuando se te solicite, introduce tu token para iniciar sesión:

>>> from huggingface_hub import notebook_login

>>> notebook_login()

Cargar el conjunto de datos Food-101

Comience cargando un subconjunto más pequeño del conjunto de datos Food-101 desde la biblioteca 🤗 Datasets. Esto le dará la oportunidad de experimentar y asegurarse de que todo funcione correctamente antes de dedicar más tiempo a entrenar con el conjunto de datos completo.

>>> from datasets import load_dataset

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

Divide la partición train del conjunto de datos en un conjunto de entrenamiento y uno de prueba con el método train_test_split:

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

A continuación, echa un vistazo a un ejemplo:

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

Cada ejemplo del conjunto de datos tiene dos campos:

• image: una imagen PIL del elemento de comida
• label: la clase de etiqueta del elemento de comida

Para facilitar que el modelo obtenga el nombre de la etiqueta a partir del id de la etiqueta, cree un diccionario que mapee el nombre de la etiqueta a un entero y viceversa:

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

Ahora puede convertir el identificador de la etiqueta en un nombre de etiqueta:

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

Preprocesar

El siguiente paso es cargar un procesador de imágenes ViT para procesar la imagen en un tensor:

>>> from transformers import AutoImageProcessor

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

Aplica algunas transformaciones de imagen a las imágenes para hacer que el modelo sea más robusto frente al sobreajuste. Aquí se utilizará el módulo transforms de torchvision, pero también puedes usar cualquier biblioteca de imágenes que prefieras.

Recorta una parte aleatoria de la imagen, redimensiona y normaliza con la media y la desviación estándar de la imagen:

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

A continuación, crea una función de preprocesamiento para aplicar las transformaciones y devolver los pixel_values de la imagen, que son las entradas del modelo:

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

Para aplicar la función de preprocesamiento a todo el conjunto de datos, utiliza el método with_transform de 🤗 Datasets. Las transformaciones se aplican al vuelo al cargar un elemento del conjunto de datos:

>>> food = food.with_transform(transforms)

Ahora crea un lote de ejemplos usando DefaultDataCollator. A diferencia de otros data collators en 🤗 Transformers, DefaultDataCollator no aplica preprocesamiento adicional, como el padding.

>>> from transformers import DefaultDataCollator

>>> data_collator = DefaultDataCollator()

Evaluar

Incluir una métrica durante el entrenamiento suele ser útil para evaluar el rendimiento de tu modelo. Puedes cargar rápidamente un método de evaluación con la biblioteca 🤗 Evaluate. Para esta tarea, carga la métrica accuracy (consulta el quick tour de 🤗 Evaluate para saber más sobre cómo cargar y calcular una métrica):

>>> import evaluate

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

Luego, crea una función que pase tus predicciones y etiquetas a compute para calcular la precisión:

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

Tu función compute_metrics ya está lista y volverás a ella cuando configures el entrenamiento.

Entrenar

Si no está familiarizado con el ajuste fino de un modelo con Trainer, ¡eche un vistazo al tutorial básico aquí!

¡Ya está listo para comenzar a entrenar su modelo! Cargue ViT con AutoModelForImageClassification. Especifique el número de etiquetas junto con el número de etiquetas esperadas y los mapeos de etiquetas:

>>> from transformers import AutoModelForImageClassification, TrainingArguments, Trainer

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

En este punto, solo quedan tres pasos:

1. Define los hiperparámetros de entrenamiento en TrainingArguments. Es importante que no elimines las columnas no utilizadas, ya que eso eliminaría la columna image. Sin la columna image, no puedes crear pixel_values. Establece remove_unused_columns=False para evitar este comportamiento. El único otro parámetro obligatorio es output_dir, que especifica dónde guardar tu modelo. Subirás este modelo al Hub estableciendo push_to_hub=True (necesitas iniciar sesión en Hugging Face para subir tu modelo). Al final de cada época, Trainer evaluará la precisión y guardará el punto de control de entrenamiento.
2. Pasa los argumentos de entrenamiento a Trainer junto con el modelo, el conjunto de datos, el tokenizador, el data collator y la función compute_metrics.
3. Llama a train() para ajustar tu 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()

Una vez completado el entrenamiento, comparte tu modelo en el Hub con el método push_to_hub() para que todo el mundo pueda usarlo:

>>> trainer.push_to_hub()

Para un ejemplo más detallado sobre cómo ajustar un modelo para la clasificación de imágenes, consulte el cuaderno de PyTorch correspondiente.

Inferencia

Genial, ahora que has ajustado un modelo, ¡puedes usarlo para inferencia!

Cargue una imagen en la que desee ejecutar inferencia:

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

La forma más sencilla de probar tu modelo afinado para inferencia es usarlo en un pipeline(). Crea una instancia de pipeline para clasificación de imágenes con tu modelo y pásale tu imagen:

>>> 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'}]

También puedes replicar manualmente los resultados de pipeline si lo deseas:

Cargue un procesador de imágenes para preprocesar la imagen y devolver la input como tensores de PyTorch:

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

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

Pasa tus entradas al modelo y devuelve los logits:

>>> from transformers import AutoModelForImageClassification

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

Obtén la etiqueta predicha con la mayor probabilidad y usa el mapeo id2label del modelo para convertirla en una etiqueta:

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