Extracción de características de imagen

La extracción de características de imágenes consiste en obtener características con significado semántico a partir de una imagen. Esta tarea tiene muchos casos de uso, como la similitud de imágenes y la recuperación de imágenes. Además, la mayoría de los modelos de visión por computadora pueden utilizarse para la extracción de características de imágenes, ya que se puede eliminar la cabeza específica de la tarea (clasificación de imágenes, detección de objetos, etc.) y obtener las características. Estas características son muy útiles a un nivel más alto: detección de bordes, detección de esquinas, entre otras. También pueden contener información sobre el mundo real (por ejemplo, cómo se ve un gato), dependiendo de la profundidad del modelo. Por lo tanto, estas salidas pueden utilizarse para entrenar nuevos clasificadores en un conjunto de datos específico.

En esta guía, vas a:

• Aprende a construir un sistema simple de similitud de imágenes sobre la base de la pipeline image-feature-extraction.
• Realiza la misma tarea con inferencia del modelo sin procesar.

Similitud de imágenes usando el Pipeline de extracción de características de imagen

Tenemos dos imágenes de gatos sentados sobre redes de pesca, una de ellas es generada.

from PIL import Image
import requests

img_urls = ["https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.png", "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats.jpeg"]
image_real = Image.open(requests.get(img_urls[0], stream=True).raw).convert("RGB")
image_gen = Image.open(requests.get(img_urls[1], stream=True).raw).convert("RGB")

Veamos el pipeline en acción. Primero, inicialice el pipeline. Si no le pasa ningún modelo, el pipeline se inicializará automáticamente con google/vit-base-patch16-224. Si desea calcular la similitud, establezca pool en True.

import torch
from transformers import pipeline
from accelerate import Accelerator
# automatically detects the underlying device type (CUDA, CPU, XPU, MPS, etc.)
device = Accelerator().device
pipe = pipeline(task="image-feature-extraction", model="google/vit-base-patch16-384", device=device, pool=True)

Para inferir con pipe, pásale ambas imágenes.

outputs = pipe([image_real, image_gen])

La salida contiene embeddings agrupados de esas dos imágenes.

# get the length of a single output
print(len(outputs[0][0]))
# show outputs
print(outputs)

# 768
# [[[-0.03909236937761307, 0.43381670117378235, -0.06913255900144577,

Para obtener la puntuación de similitud, debemos pasarlos a una función de similitud.

from torch.nn.functional import cosine_similarity

similarity_score = cosine_similarity(torch.Tensor(outputs[0]),
                                     torch.Tensor(outputs[1]), dim=1)

print(similarity_score)

# tensor([0.6043])

Si desea obtener los últimos estados ocultos antes del pooling, evite pasar ningún valor para el parámetro pool, ya que está establecido en False de forma predeterminada. Estos estados ocultos son útiles para entrenar nuevos clasificadores o modelos basados en las características del modelo.

pipe = pipeline(task="image-feature-extraction", model="google/vit-base-patch16-224", device=device)
outputs = pipe(image_real)

Dado que las salidas no están agrupadas, obtenemos los últimos estados ocultos, donde la primera dimensión es el tamaño del lote y las dos últimas son la forma de la incrustación.

import numpy as np
print(np.array(outputs).shape)
# (1, 197, 768)

Obtención de características y similitudes con AutoModel

También podemos usar la clase AutoModel de transformers para obtener las características. AutoModel carga cualquier modelo de transformers sin una cabeza específica de tarea, y podemos usarlo para obtener las características.

from transformers import AutoImageProcessor, AutoModel

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(device)

Escribamos una función simple para inferencia. Primero pasaremos las entradas al processor y luego pasaremos sus salidas al model.

def infer(image):
  inputs = processor(image, return_tensors="pt").to(device)
  outputs = model(**inputs)
  return outputs.pooler_output

Podemos pasar las imágenes directamente a esta función y obtener los embeddings.

embed_real = infer(image_real)
embed_gen = infer(image_gen)

Podemos obtener la similitud nuevamente sobre los embeddings.

from torch.nn.functional import cosine_similarity

similarity_score = cosine_similarity(embed_real, embed_gen, dim=1)
print(similarity_score)

# tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)