マルチモーダルRAG：TwelveLabsとChromaを使用した動画とのチャット

# Twelve LabsとChromaのライブラリをインストール
!pip install --upgrade twelvelabs
!pip install --upgrade chromadb

# オープンソースモデルで使用するライブラリをインストール
!pip install protobuf==3.20.3
!pip install --upgrade -q accelerate bitsandbytes
!pip install git+https://github.com/huggingface/transformers.git
!pip install av

# Colab環境以外の場合に追加でインストールするもの
# !python -m pip install pillow
# !python -m pip install sentencepiece
# !python -m pip install matplotlib

from google.colab import drive
drive.mount('/content/drive')

video_folder_path = "/content/drive/MyDrive/TwelveLabs-Chroma/videos/"

import numpy as np
import subprocess
import os

def upscale_video(input_file, output_path, target_width=854, target_height=480):

    output_file = os.path.join(output_path, os.path.basename(input_file))

    if os.path.exists(output_file):
        print(f"Skipping {input_file} as {output_file} already exists.")
        return
 

    """
    FFmpegを使用して、動画をターゲットの幅と高さにアップスケールします。

    Args:
        input_file (str): 入力動画ファイルへのパス。
        output_file (str): アップスケールされた動画を保存するパス。
        target_width (int): 希望する出力幅。デフォルトは 854。
        target_height (int): 希望する出力高さ。デフォルトは 480。
    """
    # 動画をアップスケールするためのFFmpegコマンド
    ffmpeg_command = [
        'ffmpeg',
        '-i', input_file,                              # 入力ファイル
        '-vf', f'scale={target_width}:{target_height}', # ターゲットの寸法を持つスケールフィルター
        '-c:a', 'copy',                                # 再エンコードなしでオーディオストリームをコピー
        output_file,                                    # 出力ファイル
        "-y"
    ]

    # FFmpegコマンドを実行
    subprocess.run(ffmpeg_command)

    print(f"Upscaled video saved to {output_file}")

upscaled_video_dir = video_folder_path + "upscaled_videos/"

# すべての.mp4動画をアップスケール
# 出力ディレクトリが存在しない場合は作成
if not os.path.exists(upscaled_video_dir):
    os.makedirs(upscaled_video_dir)

# 生動画ディレクトリ内の全ファイルを反復処理
for filename in os.listdir(video_folder_path):
    # ファイルが動画ファイルであるかチェック
    input_filepath = os.path.join(video_folder_path, filename)
    if filename.endswith(".mp4"):
        upscale_video(input_filepath, upscaled_video_dir)

from google.colab import userdata
TL_API_KEY=userdata.get('TL_API_KEY')

from twelvelabs import TwelveLabs
from twelvelabs.models.embed import EmbeddingsTask

# Twelve Labsクライアントを初期化
twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)

import chromadb

# 一時的なChromaクライアントを初期化
chroma_client = chromadb.Client()

# 永続クライアントのオプション
chroma_client = chromadb.PersistentClient(path="/path/to/save/to")

# HTTPクライアント
chroma_client = chromadb.HttpClient(host='localhost', port=8000)

def on_task_update(task: EmbeddingsTask):
    print(f"  Status={task.status}")

# 動画の埋め込みを作成し、Chroma向けにフォーマットします
def create_video_embeddings(client,video_file,segment_length,task_id=None):

    # Twelve Labsに動画が存在しない場合はアップロードします
    video_name = os.path.basename(video_file)

    if task_id == None or task_id == "":
        task = client.embed.task.create(
            engine_name="Marengo-retrieval-2.7",
            video_file=video_file,
            video_clip_length=segment_length
        )
        print(
            f"Created task: id={task.id} engine_name={task.engine_name} status={task.status}"
        )

        status = task.wait_for_done(
            sleep_interval=2,
            callback=on_task_update
        )

        print(f"Embedding done: {status}")

        task_id = task.id

    # 埋め込みを取得します
    task = client.embed.task.retrieve(task_id)
    print("task",task)

    # Chroma向けにフォーマットします
    embeddings = []
    metadatas = []
    ids = []

    idx = 0

    print("embeddings",task.video_embeddings)

    if task.video_embeddings is not None:
        for v in task.video_embeddings:

            metadata = {
                "embedding_scope":v.embedding_scope,
                "start_offset_sec":v.start_offset_sec,
                "end_offset_sec":v.end_offset_sec,
                "video_file":video_file,
                "video_name":video_name,
                "task_id":task.id,
                "video_segment_number":idx
            }


            embedding = v.values
            id = task.id + "_" + str(idx)

            metadatas.append(metadata)
            embeddings.append(embedding)
            ids.append(id)

            idx += 1

    return (ids,metadatas,embeddings,task_id)

# セグメントの長さと、処理対象の動画を設定します
segment_duration = 6
current_video_path = upscaled_video_dir + "How To Make Birria Tacos [4nIFJFgH99w].mp4"

# Chromaにアップロードする埋め込みを取得します

# すでにtask_idをお持ちの場合は設定し、そうでない場合は空文字を設定してください
task_id = ""
ids, metadatas, embeddings, task_id = create_video_embeddings(twelvelabs_client,current_video_path,segment_duration,task_id)

# Chromaコレクションを取得、または作成します
chroma_collection_name = "video_embeddings"
collection = chroma_client.get_or_create_collection(chroma_collection_name)

# 埋め込みとメタデータをコレクションに追加します
collection.add(
    metadatas = metadatas,
    embeddings = embeddings,
    ids=ids
)

# 最初の埋め込みをテスト検索として使用します
test_segment_embeddings = embeddings[0]

results = collection.query(
    query_embeddings=[test_segment_embeddings],
    n_results=4
)

print("search embeddings for:",ids[0])
print("found:", results["ids"][0][0])
print("distance:",results["distances"][0][0])

# 最初の動画のテキスト埋め込みがそれ自体から距離 0 であることをアサート（確認）します
assert results["ids"][0][0] == ids[0]
assert results["distances"][0][0] == 0

query = "What are the ingredients for birria tacos?"（ビリアタコスの材料は何ですか？）

import os

def query_chroma(collection,query,n_results=1):
    # クエリ用の埋め込みを作成します
    embedding = twelvelabs_client.embed.create(
        engine_name="Marengo-retrieval-2.7",
        text=query,
        text_truncate="start",
    )

    query_embeddings = embedding.text_embedding.float

    # クエリ埋め込みを使用してChromaデータベースを検索します

    response = collection.query(
        query_embeddings=query_embeddings,
        n_results=n_results,
    )


    return response

response = query_chroma(collection,query)

# 最も類似しているオブジェクトのプロパティと距離を出力します
print(response["ids"][0][0])
print(response["distances"][0][0])
print(response["metadatas"][0][0])

# 次のステップで使用するために、見つかった動画セグメントのパスを取得します
found_video_metadata = response["metadatas"][0][0]

split_video_dir = video_folder_path + "split_videos/"

def split_video(input_path, output_dir, segment_duration=6):

    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    filename = os.path.splitext(os.path.basename(input_path))[0]
    filetype = os.path.splitext(os.path.basename(input_path))[1]

    # 動画をセグメントに分割します
    ffmpeg_command = [
        'ffmpeg',
        '-i', input_path,             # 入力動画ファイル
        '-c', 'copy',                  # 映像と音声の両コーデックをコピー
        '-f', 'segment',               # セグメントモード
        '-segment_time', str(segment_duration),  # セグメントの長さ
        '-reset_timestamps', '1',      # 各セグメントのタイムスタンプをリセット
        output_dir + filename + '_%03d' + filetype  # 出力ファイル名のパターン (例: output_001.mp4)
    ]

    # コマンドを実行します
    subprocess.run(ffmpeg_command)

    print("Video split into 6-second segments successfully.")

# 動画をセグメントに分割します
split_video(input_path=current_video_path, output_dir=split_video_dir, segment_duration=segment_duration)

# pegasusのインデックスを作成、または取得します
engines = [
        {
            "name": "pegasus1.2",
            "options": ["visual", "conversation"]
        }
    ]

index_name = "cooking_video_index"
indices_list = twelvelabs_client.index.list(name=index_name)

if len(indices_list) == 0:
    index = twelvelabs_client.index.create(
        name=index_name,
        engines=engines,

    )
    print(f"A new index has been created: id={index.id} name={index.name} engines={index.engines}")
else:
    index = indices_list[0]
    print(f"Index already exists: id={index.id} name={index.name} engines={index.engines}")

# 動画セグメントのファイル名を取得します
found_video_segment_number = int(found_video_metadata["video_segment_number"])
found_video_file = found_video_metadata["video_file"]
found_video_filename = os.path.splitext(os.path.basename(found_video_file))[0]
found_video_filetype = os.path.splitext(os.path.basename(found_video_file))[1]
found_video_segment_filename = found_video_filename + f"_{found_video_segment_number:03d}"

found_video_segment_path = split_video_dir + found_video_segment_filename + found_video_filetype
print(found_video_segment_path)

def upload_video_to_twelve_labs(index,video_path):

    # 動画をtwelve labsのインデックスにアップロードします
    task = twelvelabs_client.task.create(
        index_id=index.id,
        file = video_path
    )
    print(f"Task created: id={task.id} status={task.status}")

    task.wait_for_done(sleep_interval=5, callback=on_task_update)

    if task.status != "ready":
      raise RuntimeError(f"Indexing failed with status {task.status}")
    print(f"The unique identifer of your video is {task.video_id}.")

    # video id を返します
    return task.video_id

import av
def read_video_pyav(container, indices):
    '''
    PyAVデコーダーで動画をデコードします。

    Args:
        container (av.container.input.InputContainer): PyAVコンテナ。
        indices (List[int]): デコードするフレームインデックスのリスト。

    Returns:
        np.ndarray: デコードされたフレームのnumpy配列。形状は (num_frames, height, width, 3)。
    '''
    frames = []
    container.seek(0)
    start_index = indices[0]
    end_index = indices[-1]
    for i, frame in enumerate(container.decode(video=0)):
        if i > end_index:
            break
        if i >= start_index and i in indices:
            frames.append(frame)
    return np.stack([x.to_ndarray(format="rgb24") for x in frames])

def sample_video(video_path, num_samples=8):
    container = av.open(video_path)

    # 動画から均一に num_samples 枚のフレームをサンプリングします
    total_frames = container.streams.video[0].frames
    indices = np.arange(0, total_frames, total_frames / num_samples).astype(int)

    sampled_frames = read_video_pyav(container, indices)

    return sampled_frames

def process_videos_in_folder(folder_path):
    sample_info = {}

    # サポートされている動画ファイルの拡張子
    video_extensions = ('.mp4', '.avi', '.mov', '.mkv')

    for filename in os.listdir(folder_path):
        simple_video_name = os.path.splitext(os.path.basename(filename))[0]
        if filename.lower().endswith(video_extensions):
            video_path = os.path.join(folder_path, filename)
            try:
                sampled_clip = sample_video(video_path)
                sample_info[simple_video_name] = {"sampled_video": sampled_clip, "video_path" : video_path}
            except Exception as e:
                print(f"Error processing {filename}: {str(e)}")

    return sample_info

sampled_video_info = process_videos_in_folder(split_video_dir)

# Chromaクエリで見つかった動画セグメントを取得します
video_segment = sampled_video_info[found_video_segment_filename]['sampled_video']

from transformers import BitsAndBytesConfig, LlavaNextVideoForConditionalGeneration, LlavaNextVideoProcessor
import torch

quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

processor = LlavaNextVideoProcessor.from_pretrained("llava-hf/LLaVA-NeXT-Video-7B-hf")
model = LlavaNextVideoForConditionalGeneration.from_pretrained(
    "llava-hf/LLaVA-NeXT-Video-7B-hf",
    quantization_config=quantization_config,
    device_map='auto'
)

# 後ほどノートブック内で直接動画を再生するために使用します

from matplotlib import pyplot as plt
from matplotlib import animation
from IPython.display import HTML

# 形状が (frames, height, width, channels) のnumpy配列
# 確認用にランダムに1つ選択します
video = sampled_video_info[list(sampled_video_info.keys())[0]]['sampled_video']

fig = plt.figure()
im = plt.imshow(video[0,:,:,:])

plt.close() # 生成された画像を表示させないために必要です

def init():
    im.set_data(video[0,:,:,:])

def animate(i):
    im.set_data(video[i,:,:,:])
    return im

anim = animation.FuncAnimation(fig, animate, init_func=init, frames=video.shape[0],
                               interval=100)
HTML(anim.to_html5_video())

# 各「content」は辞書のリストになっており、画像/ビデオ/テキストのモダリティを追加できます
conversation = [
      {
          "role": "user",
          "content": [
              {"type": "text", "text": query},
              {"type": "video"},
              ],
      },
]

prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
prompt_len = len(prompt)

inputs = processor([prompt], videos=[video_segment], padding=True, return_tensors="pt").to(model.device)
generate_kwargs = {"max_new_tokens": 100, "do_sample": True, "top_p": 0.9}

output = model.generate(**inputs, **generate_kwargs)
open_source_segment_generated_text = processor.batch_decode(output, skip_special_tokens=True)

print(open_source_segment_generated_text[0])

print(f"query {query}")
print("pegasus 回答")
print(segment_answer)
print("オープンソース 回答")
print(open_source_segment_generated_text[0])

# すべての動画の埋め込みを行い、タスクIDを保存します
chroma_collection_name = "video_embeddings"
collection = chroma_client.get_or_create_collection(chroma_collection_name)

# 各動画のtwelve labsのタスクIDを保存します
task_ids = {}

# 各動画の埋め込みとメタデータを取得します
# 同じ動画を複数回アップロードしないように、タスクIDを保存します
for filename in os.listdir(upscaled_video_dir):

    if filename.endswith(".mp4"):

        if (filename in task_ids.keys()):
            task_id = task_ids[filename]
        else:
            task_id = None

        file_path = os.path.join(upscaled_video_dir, filename)

        ids, metadatas, embeddings, task_id = create_video_embeddings(twelvelabs_client,file_path,segment_duration,task_id)

        task_ids[filename] = task_id

        collection.add(
            metadatas = metadatas,
            embeddings = embeddings,
            ids=ids
        )

print(task_ids)

response = query_chroma(collection,query)
found_full_video_name = response["metadatas"][0][0]["video_name"]
print(found_full_video_name)

# 同じ動画を複数回アップロードしないように、pegasusのビデオIDを保存する辞書を用意します
pegasus_video_ids = {}

for upscaled_video in os.listdir(upscaled_video_dir):
    upscaled_video_path = os.path.join(upscaled_video_dir, upscaled_video)
    print(upscaled_video_path)
    if upscaled_video not in pegasus_video_ids:
        video_id = upload_video_to_twelve_labs(index,upscaled_video_path)
        pegasus_video_ids[upscaled_video] = video_id

print(pegasus_video_ids)

video_id = pegasus_video_ids[found_full_video_name]
print(video_id)

res = twelvelabs_client.generate.text(
  video_id=video_id,
  prompt=query
)
full_video_answer = res.data
print(f"query {query}")
print(f"{full_video_answer}")

print(f"セグメントでの回答: \n{segment_answer}")

# すべての動画をサンプリングします:
sampled_database_video_info = process_videos_in_folder(upscaled_video_dir)

# 各「content」は辞書のリストになっており、画像/ビデオ/テキストのモダリティを追加できます
conversation = [
      {
          "role": "user",
          "content": [
              {"type": "text", "text": query},
              {"type": "video"},
              ],
      },
]

prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
prompt_len = len(prompt)

inputs = processor([prompt], videos=[video_segment], padding=True, return_tensors="pt").to(model.device)
generate_kwargs = {"max_new_tokens": 100, "do_sample": True, "top_p": 0.9}

output = model.generate(**inputs, **generate_kwargs)
generated_text = processor.batch_decode(output, skip_special_tokens=True)

print(generated_text[0])

print(f"Pegasusの回答: \n{full_video_answer}")