Поиск видео на YouTube с использованием естественного языка

Вы можете использовать этот блокнот для игры с нейронной сетью CLIP OpenAI для поиска видео на YouTube с использованием естественного языка.

Вы можете использовать этот блокнот для игры с нейронной сетью OpenAI CLIP для поиска видео на YouTube с использованием естественного языка.
Как это работает

Скачать видео с YouTube
Извлеките каждый N-й кадр
Кодировать все кадры, используя зажим
Кодирование поискового запроса на естественном языке с помощью CLIP
Найдите изображения, которые лучше всего соответствуют поисковому запросу

Выберите видео на YouTube

Вставьте ссылку на видео на YouTube или выберите один из примеров.

Выберите значение N, которое определяет, сколько кадров должно быть пропущено. N = 30 обычно составляет около 1 кадра в секунду.

# Dashcam video of driving around San Francisco

video_url = "https://www.youtube.com/watch?v=PGMu_Z89Ao8"  

# How much frames to skip

N = 120

Настройка окружающей среды

Установите зажим и установите его зависимостей

# pytube is used to download videos from YouTube

!pip install pytube

# Intall a newer version of plotly

!pip install plotly==4.14.3

# Install CLIP from the GitHub repo

!pip install git+https://github.com/openai/CLIP.git

# Install torch 1.7.1 with GPU support

!pip install torch==1.7.1+cu101 torchvision==0.8.2+cu101 -f https://download.pytorch.org/whl/torch_stable.html

Скачать видео и обработать с помощью КЛИПА

В этом разделе видео загружается, каждый N-й кадр извлекается и каждый кадр обрабатывается с помощью клипа.

Скачайте видео locallyat 360p resoultion с помощью pytube. Большее разрешение не требуется, потому что КЛИП все равно масштабирует изображения.

from pytube import YouTube

# Choose a video stream with resolution of 360p

streams = YouTube(video_url).streams.filter(adaptive=True, subtype="mp4", resolution="360p", only_video=True)

# Check if there is a valid stream

if len(streams) == 0:

  raise "No suitable stream found for this YouTube video!"

# Download the video as video.mp4

print("Downloading...")

streams[0].download(filename="video")

print("Download completed.")

Извлеките каждый N-й кадр видео.

import cv2

from PIL import Image

# The frame images will be stored in video_frames

video_frames = []

# Open the video file

capture = cv2.VideoCapture('video.mp4')

current_frame = 0

while capture.isOpened():

  # Read the current frame

  ret, frame = capture.read()

  # Convert it to a PIL image (required for CLIP) and store it

  if ret == True:

    video_frames.append(Image.fromarray(frame[:, :, ::-1]))

  else:

    break

  # Skip N frames

  current_frame += N

  capture.set(cv2.CAP_PROP_POS_FRAMES, current_frame)

# Print some statistics

print(f"Frames extracted: {len(video_frames)}")

Загрузите общедоступную модель клипа.

import clip

import torch

# Load the open CLIP model

device = "cuda" if torch.cuda.is_available() else "cpu"

model, preprocess = clip.load("ViT-B/32", device=device)

Кодируйте все кадры с помощью клипа. Кодирование выполняется пакетами для повышения эффективности.

import math

import numpy as np

import torch

# You can try tuning the batch size for very large videos, but it should usually be OK

batch_size = 256

batches = math.ceil(len(video_frames) / batch_size)

# The encoded features will bs stored in video_features

video_features = torch.empty([0, 512], dtype=torch.float16).to(device)

# Process each batch

for i in range(batches):

  print(f"Processing batch {i+1}/{batches}")

  # Get the relevant frames

  batch_frames = video_frames[i*batch_size : (i+1)*batch_size]

  # Preprocess the images for the batch

  batch_preprocessed = torch.stack([preprocess(frame) for frame in batch_frames]).to(device)

  # Encode with CLIP and normalize

  with torch.no_grad():

    batch_features = model.encode_image(batch_preprocessed)

    batch_features /= batch_features.norm(dim=-1, keepdim=True)

  # Append the batch to the list containing all features

  video_features = torch.cat((video_features, batch_features))

# Print some stats

print(f"Features: {video_features.shape}")

Определение функций для поиска видео

В этом разделе определяются функции, используемые для поиска видео.

import plotly.express as px

def search_video(search_query, display_heatmap=True, display_results_count=3):

  # Encode and normalize the search query using CLIP

  with torch.no_grad():

    text_features = model.encode_text(clip.tokenize(search_query).to(device))

    text_features /= text_features.norm(dim=-1, keepdim=True)

  # Compute the similarity between the search query and each frame using the Cosine similarity

  similarities = (100.0 * video_features @ text_features.T)

  values, best_photo_idx = similarities.topk(display_results_count, dim=0)

  # Display the heatmap

  if display_heatmap:

    print("Search query heatmap over the frames of the video:")

    fig = px.imshow(similarities.T.cpu().numpy(), height=50, aspect='auto', color_continuous_scale='viridis')

    fig.update_layout(coloraxis_showscale=False)

    fig.update_xaxes(showticklabels=False)

    fig.update_yaxes(showticklabels=False)

    fig.update_layout(margin=dict(l=0, r=0, b=0, t=0))

    fig.show()

    print()

  # Display the top 3 frames

  for frame_id in best_photo_idx:

    display(video_frames[frame_id])

    print()

Поисковые запросы

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

Поверх результата eah вы также найдете тепловую карту, которая показывает, насколько вероятен поисковый запрос на всех кадрах из видео.
- Пожарная машина"

search_video("a fire truck")

Поисковый запрос тепловая карта по кадрам видео:

Источник: colab.research.google.com

Комментарии: