Извлечение признаков изображений

Извлечение признаков изображений — это задача по получению семантически значимых признаков на основе изображения. У неё множество применений, включая оценку сходства изображений и поиск по изображениям. Более того, большинство моделей компьютерного зрения можно использовать для извлечения признаков изображений: достаточно удалить специфичную для задачи головную часть (классификация изображений, обнаружение объектов и т. д.), чтобы получить признаки. Эти признаки очень полезны на более высоком уровне: обнаружение границ, обнаружение углов и так далее. В зависимости от глубины модели они также могут содержать информацию о реальном мире (например, как выглядит кошка). Поэтому эти выходные данные можно использовать для обучения новых классификаторов на конкретном наборе данных.

В этом руководстве вы:

• Научитесь создавать простую систему поиска похожих изображений на основе конвейера image-feature-extraction.
• Выполните ту же задачу с помощью инференса «голой» модели.

Определение сходства изображений с помощью конвейера image-feature-extraction

У нас есть два изображения кошек, сидящих на рыболовных сетях, одно из них сгенерировано.

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

Рассмотрим работу конвейера на практике. Сначала инициализируйте конвейер. Если вы не передадите ему модель, конвейер будет автоматически инициализирован с использованием google/vit-base-patch16-224. Если вы хотите вычислить сходство, установите pool в 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)

Для выполнения инференса с помощью pipe передайте в него оба изображения.

outputs = pipe([image_real, image_gen])

Вывод содержит объединённые эмбеддинги этих двух изображений.

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

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

Чтобы получить оценку сходства, необходимо передать их в функцию сходства.

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

Если вам нужны последние скрытые состояния до пулинга, не передавайте значение параметра pool, так как по умолчанию он установлен в False. Эти скрытые состояния полезны для обучения новых классификаторов или моделей на основе признаков, извлекаемых моделью.

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

Поскольку выходные данные не пулуются, мы получаем последние скрытые состояния, где первое измерение соответствует размеру батча, а последние два — форме эмбеддинга.

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

Получение признаков и сходств с помощью AutoModel

Также можно использовать класс AutoModel из transformers для получения признаков. AutoModel загружает любую модель transformers без специализированной для конкретной задачи головы, и мы можем использовать это для извлечения признаков.

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)

Напишем простую функцию для инференса. Сначала передадим входные данные в processor, а затем его выходные данные — в model.

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

Мы можем передать изображения непосредственно в эту функцию и получить эмбеддинги.

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

Мы снова можем вычислить сходство по эмбеддингам.

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