Creación de Descripciones Automáticas de Imágenes con Transformers y Captioning, en español

En este post, exploraremos cómo utilizar modelos basados en Transformers para la generación automática de descripciones de imágenes. Implementaremos un sistema que convierte imágenes en texto mediante el modelo preentrenado VisionEncoderDecoderModel, el cual utiliza Vision Transformer (ViT) para la codificación visual y GPT-2 para la generación del lenguaje. Además, traduciremos las descripciones generadas al español usando la API de Google Translate.

¿Qué es el Image Captioning?

El image captioning es el proceso de generar una descripción textual a partir de una imagen. Esto combina dos áreas clave de la inteligencia artificial: la visión por computadora y el procesamiento del lenguaje natural (NLP). En este caso, nos apoyaremos en el modelo VisionEncoderDecoderModel de Hugging Face, que está preentrenado para realizar esta tarea de forma precisa.

Requerimientos

Antes de iniciar, asegúrate de tener instaladas las siguientes bibliotecas:

				
					pip install transformers torch pillow googletrans==4.0.0-rc1
				
			

Estas bibliotecas nos permitirán acceder al modelo de Transformers, manejar las imágenes y realizar la traducción de texto. Ahora que tenemos todo listo, pasemos al código.

Paso 1: Cargar el Modelo y Configurar el Entorno

El primer paso consiste en cargar el modelo preentrenado y configurar los componentes que utilizaremos:

				
					import torch
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from PIL import Image
from googletrans import Translator

# Cargar el modelo preentrenado y el tokenizer
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
featureExtractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

# Configurar el dispositivo (GPU o CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
				
			

En este bloque de código:

  • VisionEncoderDecoderModel: Cargamos el modelo que combina un codificador de imágenes y un generador de texto.
  • ViTImageProcessor: Preparamos el procesador de imágenes que convierte las imágenes en tensores para el modelo.
  • AutoTokenizer: Utilizamos un tokenizer que transformará los textos generados en secuencias legibles.

El modelo se cargará en la GPU si está disponible, lo que acelerará el procesamiento.

Paso 2: Definir los Parámetros de Generación y Beam Search

El siguiente paso es configurar los parámetros que controlan la forma en que el modelo generará las descripciones de las imágenes. Esto incluye la longitud máxima del texto generado y el uso de un método conocido como búsqueda por haces (beam search).

				
					# Parámetros de generación de texto
maxLength = 16
numBeams = 4
genKwargs = {"max_length": maxLength, "num_beams": numBeams}
				
			

Explicación de los Parámetros:

  • maxLength: Este valor define la longitud máxima de la descripción generada. En este caso, limitamos la salida a un máximo de 16 tokens, es decir, 16 palabras o caracteres que formarán la oración resultante. Si la descripción generada es más corta, el modelo dejará de generar texto antes de alcanzar este límite.
  • numBeams: Este parámetro es especialmente importante y se refiere al número de “haces” utilizados en el proceso de generación de texto. Beam search o búsqueda por haces es un método de búsqueda heurística que permite al modelo considerar múltiples alternativas de palabras en cada paso de la generación.

Paso 3: Función de Predicción de Captions

Ahora definimos una función para procesar las imágenes y generar las captions. Esta función tomará una lista de rutas de imágenes, las convertirá en tensores y las pasará al modelo para generar las descripciones.

				
					def predictStep(imagePaths):
    images = []
    for imagePath in imagePaths:
        iImage = Image.open(imagePath)
        if iImage.mode != "RGB":
            iImage = iImage.convert(mode="RGB")
        images.append(iImage)

    pixelValues = featureExtractor(images=images, return_tensors="pt").pixel_values
    pixelValues = pixelValues.to(device)

    outputIds = model.generate(pixelValues, **genKwargs)
    preds = tokenizer.batch_decode(outputIds, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds
				
			

En esta función:

  • Image.open: Cargamos la imagen desde la ruta proporcionada y la convertimos a formato RGB si es necesario.
  • featureExtractor: Procesa la imagen y la convierte en un tensor que el modelo pueda entender.
  • model.generate: Genera la caption usando el modelo preentrenado con los parámetros definidos.
  • tokenizer.batch_decode: Decodifica la salida del modelo para convertirla en texto.

Paso 4: Traducción de las Captions al Español

Una vez generadas las descripciones, podemos traducirlas del inglés al español con Google Translate.

				
					# Función para traducir captions al español
def translateCaption(caption):
    translator = Translator()
    translation = translator.translate(caption, dest="es", src="en")
    return translation.text
				
			

Aquí, usamos la API de Google Translate para traducir el texto generado del inglés al español.

Paso 5: Probar el Sistema Completo

Finalmente, podemos probar todo el sistema generando una descripción para una imagen y traducirla.

				
					# Generar y traducir la caption de una imagen
imageCaption = predictStep(['0.jpg'])[0]
translatedCaption = translateCaption(imageCaption)

print(f"Caption en inglés: {imageCaption}")
print(f"Caption traducida al español: {translatedCaption}")
				
			

Aquí tienes el código completo ordenado de una mejor manera.

				
					import torch
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from PIL import Image
from googletrans import Translator


def predictStep(imagePaths):
    images = []
    for imagePath in imagePaths:
        iImage = Image.open(imagePath)
        if iImage.mode != "RGB":
            iImage = iImage.convert(mode="RGB")
        images.append(iImage)

    pixelValues = featureExtractor(images=images, return_tensors="pt").pixel_values
    pixelValues = pixelValues.to(device)

    outputIds = model.generate(pixelValues, **genKwargs)
    preds = tokenizer.batch_decode(outputIds, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds


# Función para traducir captions al español
def translateCaption(caption):
    translator = Translator()
    translation = translator.translate(caption, dest="es", src="en")
    return translation.text

# Cargar el modelo preentrenado y el tokenizer
model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
featureExtractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

# Configurar el dispositivo (GPU o CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Parámetros de generación de texto
maxLength = 16
numBeams = 4
genKwargs = {"max_length": maxLength, "num_beams": numBeams}

# Generar y traducir la caption de una imagen
imageCaption = predictStep(['0.jpg'])[0]
translatedCaption = translateCaption(imageCaption)

print(f"Caption en inglés: {imageCaption}")
print(f"Caption traducida al español: {translatedCaption}")
				
			

Ejemplo de Resultados

Probemos el sistema con una imagen llamada 1.jpg. El modelo generará una descripción y luego la traducirá al español. Un posible resultado sería:

Respuestas