ビデオの干し草の山から針を探す：TwelveLabs Marengo、Amazon Bedrock、Elasticsearchによるインテリジェントなビデオ検索の構築

import os, json, time, copy
import boto3, botocore
import yt_dlp
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk
from pathlib import Path
from dotenv import load_dotenv
import tqdm

# Configuration
AWS_REGION = "us-east-1"
S3_BUCKET_NAME = os.getenv("S3_BUCKET_NAME")
MARENGO_MODEL_ID = "twelvelabs.marengo-embed-2-7-v1:0"  
TEXT_EMBEDDING_MODEL_ID = "us.twelvelabs.marengo-embed-2-7-v1:0"
ELASTICSEARCH_ENDPOINT = os.getenv("ELASTICSEARCH_ENDPOINT")
ELASTICSEARCH_API_KEY = os.getenv("ELASTICSEARCH_API_KEY")

# Test dataset: 2025 summer blockbuster trailers
videos = [
    "https://www.youtube.com/watch?v=VWqJifMMgZE",  # Lilo and Stitch 2025
    "https://www.youtube.com/watch?v=Ox8ZLF6cGM0",  # Superman 2025
    "https://www.youtube.com/watch?v=jan5CFWs9ic",  # Jurassic World Rebirth
    "https://www.youtube.com/watch?v=qpoBjOg5RHU",  # Fantastic Four: First Steps
    "https://www.youtube.com/watch?v=22w7z_lT6YM",  # How to Train Your Dragon
]

def get_video(video: VideoIntelligence):
    """Download video using yt-dlp and extract metadata"""
    base_directory = Path(video.get_base_path())
    base_directory.mkdir(parents=True, exist_ok=True)
    
    video_path = f"{video.get_base_path()}/{video.get_video_string()}.mp4"
    metadata_path = f"{video.get_base_path()}/metadata.json"
    
    ydl_opts = {
        "format": "bestvideo+bestaudio/best",
        "outtmpl": video_path,
        "merge_output_format": "mp4"
    }
    
    if not os.path.exists(video_path):
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            metadata = ydl.extract_info(video.url, download=False)
            ydl.download([video.url])
        with open(metadata_path, "w") as f:
            json.dump(metadata, f)
    else:
        metadata = json.load(open(metadata_path, "r"))
    
    video.set_metadata(metadata) 
    video.set_video_path(video_path)

# Process all videos
video_objects = []
for video_str in videos:
    if "youtube.com" in video_str:
        platform = "youtube"
        video_id = video_str.split("v=")[1] 
        video_objects.append(VideoIntelligence(video_str, platform, video_id))

for video_object in video_objects:
    get_video(video_object)

# Initialize AWS clients
session = boto3.session.Session(
    aws_access_key_id=AWS_ACCESS_KEY_ID,
    aws_secret_access_key=AWS_SECRET_ACCESS_KEY, 
    region_name=AWS_REGION
)

aws_account_id = session.client('sts').get_caller_identity()["Account"]
s3_client = session.client('s3')
bedrock_client = session.client("bedrock-runtime", region_name=AWS_REGION)

# Upload videos to S3 with existence checking
for video_object in video_objects:
    video_path = video_object.get_video_path()
    s3_key = f"videos/{video_object.get_platform()}/{video_object.get_video_id()}/{os.path.basename(video_path)}"
    video_object.set_s3_key(s3_key)
    
    try:
        s3_client.head_object(Bucket=S3_BUCKET_NAME, Key=s3_key)
        print(f"Video {video_object.get_video_string()} already exists in S3")
        continue
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == '404':
            print(f"Uploading {video_object.get_video_string()} to S3...")
            s3_client.upload_file(video_path, S3_BUCKET_NAME, s3_key)
            print(f"Successfully uploaded {video_object.get_video_string()}")

def wait_for_embedding_output(s3_bucket: str, s3_prefix: str, 
                             invocation_arn: str, verbose: bool = False) -> list:
    """Poll Bedrock async job until completion and retrieve results"""
    status = None
    while status not in ["Completed", "Failed", "Expired"]:
        response = bedrock_client.get_async_invoke(invocationArn=invocation_arn)
        status = response['status']
        if verbose:
            tqdm.tqdm.write(f"Embedding task status: {status}")
        time.sleep(5)
    
    if status != "Completed":
        raise Exception(f"Embedding task failed with status: {status}")
    
    # Retrieve output from S3
    response = s3_client.list_objects_v2(Bucket=s3_bucket, Prefix=s3_prefix)
    for obj in response.get('Contents', []):
        if obj['Key'].endswith('output.json'):
            output_key = obj['Key']
            obj = s3_client.get_object(Bucket=s3_bucket, Key=output_key)
            content = obj['Body'].read().decode('utf-8')
            return json.loads(content).get("data", [])
    
    raise Exception("No output.json found in S3 prefix")

def create_video_embedding(video_s3_uri: str, video_id: str) -> list:
    """Start async Marengo embedding job for video in S3"""
    s3_output_prefix = f'embeddings/videos/{video_id}'
    
    response = bedrock_client.start_async_invoke(
        modelId=MARENGO_MODEL_ID,
        modelInput={
            "inputType": "video",
            "mediaSource": {
                "s3Location": {"uri": video_s3_uri, "bucketOwner": aws_account_id}
            }
        },
        outputDataConfig={
            "s3OutputDataConfig": {"s3Uri": f's3://{S3_BUCKET_NAME}/{s3_output_prefix}'}
        }
    )
    
    invocation_arn = response["invocationArn"]
    print(f"Video embedding task started: {invocation_arn}")
    
    return wait_for_embedding_output(S3_BUCKET_NAME, s3_output_prefix, invocation_arn)

def check_existing_embedding(video_id: str):
    """Check S3 for cached embeddings to avoid re-processing"""
    s3_output_prefix = f'embeddings/videos/{video_id}'
    
    try:
        response = s3_client.list_objects_v2(Bucket=S3_BUCKET_NAME, Prefix=s3_output_prefix)
        if 'Contents' in response and any(obj['Key'].endswith('output.json') for obj in response.get('Contents', [])):
            # Load existing embeddings from S3
            for obj in response.get('Contents', []):
                if obj['Key'].endswith('output.json'):
                    output_key = obj['Key']
                    obj = s3_client.get_object(Bucket=S3_BUCKET_NAME, Key=output_key)
                    content = obj['Body'].read().decode('utf-8')
                    return json.loads(content).get("data", [])
        return None
    except botocore.exceptions.ClientError as e:
        if e.response['Error']['Code'] == '404':
            return None
        raise e

# Generate embeddings with caching
for video_object in tqdm.tqdm(video_objects, desc="Processing videos"):
    video_id = video_object.get_video_id()
    video_uri = f"s3://{S3_BUCKET_NAME}/{video_object.get_s3_key()}"
    
    # Check for existing embeddings first
    retrieved_embeddings = check_existing_embedding(video_id)
    if retrieved_embeddings:
        video_object.set_embeddings_list(retrieved_embeddings)
    else:
        embedding_data = create_video_embedding(video_uri, video_id)
        video_object.set_embeddings_list(embedding_data)

# Connect to Elasticsearch
es = Elasticsearch(
    hosts=[ELASTICSEARCH_ENDPOINT],
    api_key=ELASTICSEARCH_API_KEY
)

# Prepare documents for indexing
docs = []
for video_object in video_objects:
    persist_object = video_object.get_video_object()
    embeddings = video_object.get_embeddings_list()
    
    for embedding in embeddings:
        if embedding["embeddingOption"] == "visual-image":  # Filter for visual embeddings
            doc = copy.deepcopy(persist_object) 
            doc["embedding"] = embedding["embedding"]
            doc["start_sec"] = embedding["startSec"] 
            doc["end_sec"] = embedding["endSec"]
            docs.append(doc)

# Create indices for different vector search methods
index_varieties = [
    "flat",       # Brute force, highest accuracy
    "hnsw",       # Hierarchical navigable small world graph
    "int8_hnsw",  # Quantized for efficiency
    "bbq_hnsw",   # Better Binary Quantization with HNSW
    "bbq_flat"    # BBQ with flat search
]

for index_variety in index_varieties:
    index_name = f"twelvelabs-movie-trailer-{index_variety}"
    mappings = {
        "properties": {
            "url": {"type": "keyword"},
            "platform": {"type": "keyword"},
            "video_id": {"type": "keyword"},
            "title": {"type": "text", "analyzer": "standard"},
            "embedding": {
                "type": "dense_vector", 
                "dims": 1024,
                "similarity": "cosine",
                "index_options": {
                    "type": index_variety
                }
            },
            "start_sec": {"type": "float"},
            "end_sec": {"type": "float"}
        }
    }
    
    # Recreate index if exists
    if es.indices.exists(index=index_name):
        es.indices.delete(index=index_name)
        time.sleep(2)
    
    es.indices.create(index=index_name, mappings=mappings)
    print(f"Index '{index_name}' created successfully")
    
    # Bulk insert documents
    actions = [{"_index": index_name, "_source": doc} for doc in docs]
    success, failed = bulk(es, actions, chunk_size=100, max_retries=3)
    print(f"Successfully indexed {success} documents into {index_name}")

# Generate text embedding for search
def create_text_embedding(text_query: str) -> list:
    """Generate text embedding using Marengo via Bedrock"""
    text_model_input = {"inputType": "text", "inputText": text_query}
    
    response = bedrock_client.invoke_model(
        modelId=TEXT_EMBEDDING_MODEL_ID,
        body=json.dumps(text_model_input)
    )
    
    response_body = json.loads(response['body'].read().decode('utf-8'))
    embedding_data = response_body.get("data", [])
    
    return embedding_data[0]["embedding"] if embedding_data else None

def vector_query(index_name: str, text_query: str) -> dict:
    """Execute k-NN vector search against Elasticsearch"""
    query_embedding = create_text_embedding(text_query)
    
    query = {
        "retriever": {
            "knn": {
                "field": "embedding",
                "query_vector": query_embedding,
                "k": 10,
                "num_candidates": "25"  
            }
        },
        "size": 10,
        "_source": False,
        "fields": ["title", "video_id", "start_sec"]
    }
    
    return es.search(index=index_name, body=query).body

# Test search
text_query = "Show me scenes with dinosaurs"  
results = vector_query("twelvelabs-movie-trailer-flat", text_query)
print(results)

# Clean up indexes when done
for index_variety in index_varieties:
    es.indices.delete(index=f"twelvelabs-movie-trailer-{index_variety}")

# Clean up S3 objects
s3_objects = s3_client.list_objects_v2(Bucket=S3_BUCKET_NAME)
if 'Contents' in s3_objects:
    delete_keys = [{'Key': obj['Key']} for obj in s3_objects['Contents']]
    s3_client.delete_objects(Bucket=S3_BUCKET_NAME, Delete={'Objects': delete_keys})