Extração de características de imagem

A extração de características de imagens é a tarefa de extrair características semanticamente significativas a partir de uma imagem. Isso tem muitos casos de uso, incluindo similaridade de imagens e recuperação de imagens. Além disso, a maioria dos modelos de visão computacional pode ser usada para extração de características de imagens, onde é possível remover a cabeça específica da tarefa (classificação de imagens, detecção de objetos, etc.) e obter as características. Essas características são muito úteis em um nível mais alto: detecção de bordas, detecção de cantos e assim por diante. Elas também podem conter informações sobre o mundo real (por exemplo, como um gato se parece), dependendo da profundidade do modelo. Portanto, essas saídas podem ser usadas para treinar novos classificadores em um conjunto de dados específico.

Neste guia, você vai:

• Aprenda a construir um sistema simples de similaridade de imagens com base no pipeline image-feature-extraction.
• Conclua a mesma tarefa com inferência do modelo puro.

Similaridade de imagens usando o Pipeline image-feature-extraction

Temos duas imagens de gatos sentados sobre redes de pesca, uma delas é gerada.

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

Vamos ver o pipeline em ação. Primeiro, inicialize o pipeline. Se você não passar nenhum modelo, o pipeline será inicializado automaticamente com google/vit-base-patch16-224. Se você quiser calcular a similaridade, defina pool como 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 com pipe, passe as duas imagens para ele.

outputs = pipe([image_real, image_gen])

A saída contém embeddings agregados dessas duas imagens.

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

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

Para obter a pontuação de similaridade, precisamos passá-los para uma função de similaridade.

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

Se você deseja obter os últimos estados ocultos antes do pooling, evite passar qualquer valor para o parâmetro pool, pois ele é definido como False por padrão. Esses estados ocultos são úteis para treinar novos classificadores ou modelos com base nos recursos do modelo.

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

Como as saídas não são agrupadas, obtemos os últimos estados ocultos, em que a primeira dimensão é o tamanho do lote e as duas últimas são a forma do embedding.

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

Obtendo recursos e similaridades com AutoModel

Também podemos usar a classe AutoModel do transformers para obter os recursos. O AutoModel carrega qualquer modelo do transformers sem uma cabeça específica de tarefa, e podemos usá-lo para obter os recursos.

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)

Vamos escrever uma função simples para inferência. Primeiro, passaremos as entradas para o processor e, em seguida, passaremos suas saídas para o model.

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

Podemos passar as imagens diretamente para esta função e obter os embeddings.

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

Podemos obter novamente a similaridade entre os 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>)