S3D MIL-NCE と OpenVINO によるテキストからのビデオ検索#

import platform 

%pip install -Uq pip 
%pip install --upgrade --pre openvino-tokenizers openvino --extra-index-url "https://storage.openvinotoolkit.org/simple/wheels/nightly" 
%pip install -q "tensorflow-macos>=2.5; sys_platform == 'darwin' and platform_machine == 'arm64' and python_version > '3.8'" # macOS M1 and M2 
%pip install -q "tensorflow-macos>=2.5,<=2.12.0; sys_platform == 'darwin' and platform_machine == 'arm64' and python_version <= '3.8'" # macOS M1 and M2 
%pip install -q "tensorflow>=2.5; sys_platform == 'darwin' and platform_machine != 'arm64' and python_version > '3.8'" # macOS x86 
%pip install -q "tensorflow>=2.5,<=2.12.0; sys_platform == 'darwin' and platform_machine != 'arm64' and python_version <= '3.8'" # macOS x86 
%pip install -q "tensorflow>=2.5; sys_platform != 'darwin' and python_version > '3.8'" 
%pip install -q "tensorflow>=2.5,<=2.12.0; sys_platform != 'darwin' and python_version <= '3.8'" 

%pip install -q tensorflow_hub tf_keras numpy "opencv-python" "nncf>=2.10.0" 
if platform.system() != "Windows":
     %pip install -q "matplotlib>=3.4" 
else:
     %pip install -q "matplotlib>=3.4,<3.7"

import os 
from pathlib import Path 

import tensorflow as tf 
import tensorflow_hub as hub 

import numpy as np 
import cv2 
from IPython import display 
import math 

os.environ["TFHUB_CACHE_DIR"] = str(Path("./tfhub_modules").resolve())

hub_handle = "https://www.kaggle.com/models/deepmind/mil-nce/TensorFlow1/s3d/1" 
hub_model = hub.load(hub_handle)

def generate_embeddings(model, input_frames, input_words): 
    """Generate embeddings from the model from video frames and input words."""
    # Input_frames は [0, 1] で正規化され、形状はバッチ x T x H x W x 3 である必要があります 
    vision_output = model.signatures["video"](tf.constant(tf.cast(input_frames, dtype=tf.float32))) 
    text_output = model.signatures["text"](tf.constant(input_words)) 

    return vision_output["video_embedding"], text_output["text_embedding"]

# @title Define video loading and visualization functions { display-mode: "form" } 

# CV2 を使用してビデオファイルを開くユーティリティー 
def crop_center_square(frame): 
    y, x = frame.shape[0:2] 
    min_dim = min(y, x) 
    start_x = (x // 2) - (min_dim // 2) 
    start_y = (y // 2) - (min_dim // 2) 
    return frame[start_y : start_y + min_dim, start_x : start_x + min_dim] 

def load_video(video_url, max_frames=32, resize=(224, 224)): 
    path = tf.keras.utils.get_file(os.path.basename(video_url)[-128:], video_url) 
    cap = cv2.VideoCapture(path) 
    frames = [] 
    try: 
        while True: 
            ret, frame = cap.read() 
            if not ret: 
                break 
            frame = crop_center_square(frame) 
            frame = cv2.resize(frame, resize) 
            frame = frame[:, :, [2, 1, 0]] 
            frames.append(frame) 

            if len(frames) == max_frames: 
                break 
    finally: 
        cap.release() 
    frames = np.array(frames) 
    if len(frames) < max_frames: 
        n_repeat = int(math.ceil(max_frames / float(len(frames)))) 
        frames = frames.repeat(n_repeat, axis=0) 
    frames = frames[:max_frames] 
    return frames / 255.0 

def display_video(urls): 
    html = "<table>" 
    html += "<tr><th>Video 1</th><th>Video 2</th><th>Video 3</th></tr><tr>" 
    for url in urls: 
        html += "<td>" 
        html += '<img src="{}" height="224">'.format(url) 
        html += "</td>" 
    html += "</tr></table>" 
    return display.HTML(html) 

def display_query_and_results_video(query, urls, scores): 
    """Display a text query and the top result videos and scores.""" 
    sorted_ix = np.argsort(-scores) 
    html = "" 
    html += "<h2>Input query: <i>{}</i> </h2><div>".format(query) 
    html += "Results: <div>" html += "<table>" 
    html += "<tr><th>Rank #1, Score:{:.2f}</th>".format(scores[sorted_ix[0]]) 
    html += "<th>Rank #2, Score:{:.2f}</th>".format(scores[sorted_ix[1]]) 
    html += "<th>Rank #3, Score:{:.2f}</th></tr><tr>".format(scores[sorted_ix[2]]) 
    for i, idx in enumerate(sorted_ix): 
        url = urls[sorted_ix[i]] 
        html += "<td>" 
        html += '<img src="{}" height="224">'.format(url) 
        html += "</td>" 
    html += "</tr></table>" 
    return html

# @title Load example videos and define text queries { display-mode: "form" } 

video_1_url = "https://upload.wikimedia.org/wikipedia/commons/b/b0/YosriAirTerjun.gif" # @param {type:"string"} 
video_2_url = "https://upload.wikimedia.org/wikipedia/commons/e/e6/Guitar_solo_gif.gif" # @param {type:"string"} 
video_3_url = "https://upload.wikimedia.org/wikipedia/commons/3/30/2009-08-16-autodrift-by-RalfR-gif-by-wau.gif" # @param {type:"string"} 

video_1 = load_video(video_1_url) 
video_2 = load_video(video_2_url) 
video_3 = load_video(video_3_url) 
all_videos = [video_1, video_2, video_3] 

query_1_video = "waterfall" # @param {type:"string"} 
query_2_video = "playing guitar" # @param {type:"string"} 
query_3_video = "car drifting" # @param {type:"string"} 
all_queries_video = [query_1_video, query_2_video, query_3_video] 
all_videos_urls = [video_1_url, video_2_url, video_3_url] 
display_video(all_videos_urls)

# ビデオ入力を準備 
videos_np = np.stack(all_videos, axis=0) 

# テキスト入力を準備 
words_np = np.array(all_queries_video) 

# ビデオとテキストの埋め込みを生成 
video_embd, text_embd = generate_embeddings(hub_model, videos_np, words_np) 

# ビデオとテキスト間のスコアはドット積によって計算されます 
all_scores = np.dot(text_embd, tf.transpose(video_embd))

# 結果を表示 
html = "" 
for i, words in enumerate(words_np): 
    html += display_query_and_results_video(words, all_videos_urls, all_scores[i, :]) 
    html += "<br>" 
display.HTML(html)

import openvino_tokenizers # NOQA Need to import conversion and operation extensions 
import openvino as ov 

model_path = hub.resolve(hub_handle) 
# ランダムデータから推論 
images_data = np.random.rand(3, 32, 224, 224, 3).astype(np.float32) 
words_data = np.array(["First sentence", "Second one", "Abracadabra"], dtype=str) 

ov_model = ov.convert_model(model_path, input=[("words", [3]), ("images", [3, 32, 224, 224, 3])])

core = ov.Core() 

compiled_model = core.compile_model(ov_model, device_name="CPU")

# コンパイル IR モデルを使用できるように `generate_embeddings` 関数を再定義 
def generate_embeddings(model, input_frames, input_words): 
    """Generate embeddings from the model from video frames and input words."""
    # Input_frames は [0, 1] で正規化され、形状はバッチ x T x H x W x 3 である必要があります 
    output = compiled_model({"words": input_words, "images": tf.cast(input_frames, dtype=tf.float32)}) 

    return output["video_embedding"], output["text_embedding"]

# ビデオとテキストの埋め込みを生成 
video_embd, text_embd = generate_embeddings(compiled_model, videos_np, words_np) 

# ビデオとテキスト間のスコアはドット積によって計算されます 
all_scores = np.dot(text_embd, tf.transpose(video_embd))

# 結果を表示 
html = "" 
for i, words in enumerate(words_np): 
    html += display_query_and_results_video(words, all_videos_urls, all_scores[i, :]) 
    html += "<br>" 
display.HTML(html)

int8_model = core.compile_model(quantized_model_path, device_name="CPU")

# ビデオとテキストの埋め込みを生成 
video_embd, text_embd = generate_embeddings(int8_model, videos_np, words_np) 

# ビデオとテキスト間のスコアはドット積によって計算されます 
all_scores = np.dot(text_embd, tf.transpose(video_embd))

# 結果を表示 
html = "" 
for i, words in enumerate(words_np): 
    html += display_query_and_results_video(words, all_videos_urls, all_scores[i, :]) 
    html += "<br>" 
display.HTML(html)