Мы реализуем наивный семантический поиск видео, используя модель OpenAI CLIP в Python (игнорируя звук). (Репозиторий GitHub для этого поста здесь.)

К концу поста мы сможем искать контент в видео, описывая его словами:

query = "A man hanging from a boom barrier"
frame, similarities = get_most_similar_frame(query)
display_frame(frame)
plot_search(query, similarities)

Для начала давайте импортируем нужные нам пакеты и загрузим модель (около 340 мб).

from pathlib import Path

import clip
import cv2
import matplotlib.pyplot as plt
import seaborn as sns
import torch
from PIL import Image
from tqdm import tqdm

sns.set_theme()
torch.set_printoptions(sci_mode=False)
data_dir = Path("data_dir")
device = "cuda" if torch.cuda.is_available() else "cpu"
model, preprocess = clip.load("ViT-B/32", device=device)  # will download ~340mb model

Краткий обзор того, как работает CLIP

Короче говоря, CLIP кодирует изображения и текст в одно и то же векторное пространство.

Более конкретно, CLIP может взять изображение и превратить его в вектор размером 512, а CLIP может взять фрагмент текста и превратить его в вектор размера 512. Если текст и изображение похожи, их векторы будут «похожими» (через косинусное сходство — то есть векторы должны указывать в одинаковом направлении). Вот пример:

image = preprocess(Image.open("dog.jpeg")).unsqueeze(0).to(device)
dog_text = clip.tokenize(["a photo of a dog"]).to(device)
cat_text = clip.tokenize(["a photo of a cat"]).to(device)
misc_text = clip.tokenize(["red green yellow"]).to(device)
with torch.no_grad():
    image_vector = model.encode_image(image)
    dog_text_vector = model.encode_text(dog_text)
    cat_text_vector = model.encode_text(cat_text)
    misc_text_vector = model.encode_text(misc_text)
print(f"{image_vector.shape = }")
print(f"{dog_text_vector.shape = }")
print(f"{cat_text_vector.shape = }")
print(f"{misc_text_vector.shape = }")
image_vector.shape = torch.Size([1, 512])
dog_text_vector.shape = torch.Size([1, 512])
cat_text_vector.shape = torch.Size([1, 512])
misc_text_vector.shape = torch.Size([1, 512])
dog_similarity = torch.cosine_similarity(image_vector, dog_text_vector).item()
cat_similarity = torch.cosine_similarity(image_vector, cat_text_vector).item()
misc_similarity = torch.cosine_similarity(image_vector, misc_text_vector).item()
print(f"dog similarity: {dog_similarity:.2f}")
print(f"cat similarity: {cat_similarity:.2f}")
print(f"misc similarity: {misc_similarity:.2f}")
dog similarity: 0.28
cat similarity: 0.20
misc similarity: 0.20

Преобразование изображений и текста в векторы:

Обратите внимание, что собака имеет наивысшее значение, на что мы и надеемся, поскольку на изображении изображена собака.

Но кажутся ли значения для кошек и разного достаточно низкими по сравнению со значением для собак?

Итак, глядя на кодовую базу CLIP, мы видим, что softmax с параметром температуры (т.е. logit_scale) используется для подобия косинуса, например:

logit_scale = model.logit_scale.exp().item()
distances = logit_scale * torch.tensor(
    [dog_similarity, cat_similarity, misc_similarity]
)
softmaxed_distances = distances.exp() / distances.exp().sum()
print(f"{logit_scale = }")
print(f"{softmaxed_distances = }")
for x, distance in zip(["dog", "cat", "misc"], softmaxed_distances):
    print(f"{x} similarity: {distance.item():.4f}")
logit_scale = 100.0
softmaxed_distances = tensor([    0.9990,     0.0005,     0.0005])
dog similarity: 0.9990
cat similarity: 0.0005
misc similarity: 0.0005

Итак, мы видим, что модель вполне уверена, что «фото собаки» — лучший из предложенных ей вариантов описания изображения.

Обратите внимание, что на самом деле нам не нужно делать эти вещи с softmax, так как мы просто заботимся о максимальном значении, которое сохраняется (например, собака все еще была максимальным значением до softmax), но это помогает сделать результаты более интерпретируемыми.

Теперь, как мы могли бы использовать эту технологию для выполнения семантического поиска по видео?…

Скачать видео

Мы возьмем видео с различным визуальным содержанием (поскольку мы игнорируем звук). (Это сборник трюков Бастера Китона, зацените.)

video_path = data_dir / 'buster_keaton.mp4'
if not video_path.is_file():
    !yt-dlp -f 133 -o {video_path}  \
        'https://www.youtube.com/watch?v=frYIj2FGmMA'

Преобразование видео в векторы CLIP

Используйте OpenCV (cv2) для обработки видео.

cap = cv2.VideoCapture(str(video_path))
frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = int(cap.get(cv2.CAP_PROP_FPS))
print(f"{frame_count = }")
print(f"{fps = }")
print(
    f"video size: {cap.get(cv2.CAP_PROP_FRAME_WIDTH):.0f}w {cap.get(cv2.CAP_PROP_FRAME_HEIGHT):.0f}h"
)
frame_count = 7465
fps = 24
video size: 426w 240h

Переберите все кадры в видео и преобразуйте их в векторы CLIP:

image_vectors = torch.zeros((frame_count, 512), device=device)
for i in tqdm(range(frame_count)):
    ret, frame = cap.read()
    with torch.no_grad():
        image_vectors[i] = model.encode_image(
            preprocess(Image.fromarray(frame)).unsqueeze(0).to(device)
        )
100%|██████████| 7465/7465 [01:17<00:00, 96.33it/s]

Большая часть работы была сделана здесь.

Далее мы создаем функцию, которая будет преобразовывать «запрос» (то есть фрагмент текста) в вектор и искать наиболее похожие видеокадры.

def get_most_similar_frame(query: str) -> tuple[int, list[float]]:
    query_vector = model.encode_text(clip.tokenize([query]).to(device))
    similarities = torch.cosine_similarity(image_vectors, query_vector)
    index = similarities.argmax().item()
    return index, similarities.squeeze()


def display_frame(index: int):
    cap.set(cv2.CAP_PROP_POS_FRAMES, index)
    ret, frame = cap.read()
    display(Image.fromarray(frame))


def plot_search(query, similarities):
    plt.figure(figsize=(8, 4))
    plt.plot((logit_scale * similarities).softmax(dim=0).tolist())
    plt.title(f"Search of video frames for '{query}'")
    plt.xlabel("Frame number")
    plt.ylabel("Query-frame similarity (softmaxed)")
    plt.show()

Полученные результаты

query = "A man sitting on the front of a train"
frame, similarities = get_most_similar_frame(query)
display_frame(frame)
plot_search(query, similarities)

query = "A woman in water"
frame, similarities = get_most_similar_frame(query)
display_frame(frame)
plot_search(query, similarities)

query = "A man on a collapsing car"
frame, similarities = get_most_similar_frame(query)
display_frame(frame)
plot_search(query, similarities)

query = "A woman and man hanging in front of a waterfall"
frame, similarities = get_most_similar_frame(query)
display_frame(frame)
plot_search(query, similarities)

Есть много возможностей для улучшения, но этот наивный подход заводит нас на удивление далеко… (ну, возможно, это и не удивительно, учитывая, что мы интенсивно используем невероятную технологию, CLIP).

Улучшения (просто выкидываю идеи):

  • Пакетное кодирование видеокадров для ускорения
  • Пропускать очень похожие видеокадры (до и/или после встраивания)
  • Используйте инструменты быстрого поиска по сходству векторов, например. файсс
  • Возможна лучшая модель встраивания (например, увеличенная модель клипа).
  • Предоставление лучших результатов K (может обеспечить некоторый уровень разнообразия/пиков использования, которые разнесены во времени)
  • И многое другое…

Это был забавный POC, и это так здорово, что такие модели, как CLIP, доступны для свободного использования!

Первоначально опубликовано на https://sidsite.com 17 ноября 2022 г.