TwelveLabsとApertureDBによるセマンティックビデオ検索エンジン

# Install necessary libraries and dependencies
!pip install requests
!pip install aperturedb
!pip install twelvelabs

# Import required modules
import requests
import aperturedb
import twelvelabs

from google.colab import userdata

# Configure Twelve Labs
TL_API_KEY = userdata.get('TL_API_KEY')

# Configure ApertureDB
ADB_PASSWORD = userdata.get('ADB_PASSWORD')

# Verify that the keys are properly set
if not TL_API_KEY:
    raise ValueError("Twelve Labs API key not found. Please set it in Colab's Secrets.")
if not ADB_PASSWORD:
    raise ValueError("ApertureDB password not found. Please set it in Colab's Secrets.")

print("API keys successfully configured.")

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

# Initialize the Twelve Labs client
twelvelabs_client = TwelveLabs(api_key=TL_API_KEY)

def generate_embedding(video_url):
    # Create an embedding task
    task = twelvelabs_client.embed.task.create(
        engine_name="Marengo-retrieval-2.6",
        video_url=video_url
    )
    print(f"Created task: id={task.id} engine_name={task.engine_name} status={task.status}")

    # Define a callback function to monitor task progress
    def on_task_update(task: EmbeddingsTask):
        print(f"  Status={task.status}")

    # Wait for the task to complete
    status = task.wait_for_done(
        sleep_interval=2,
        callback=on_task_update
    )
    print(f"Embedding done: {status}")

    # Retrieve the task result
    task_result = twelvelabs_client.embed.task.retrieve(task.id)

    # Extract and return the embeddings
    embeddings = []
    for v in task_result.video_embeddings:
        embeddings.append({
            'embedding': v.embedding.float,
            'start_offset_sec': v.start_offset_sec,
            'end_offset_sec': v.end_offset_sec,
            'embedding_scope': v.embedding_scope
        })

    return embeddings, task_result

# Example usage
video_url = "https://storage.googleapis.com/ad-demos-datasets/videos/Ecommerce%20v2.5.mp4"

# Generate embeddings for the video
embeddings, task_result = generate_embedding(video_url)

print(f"Generated {len(embeddings)} embeddings for the video")
for i, emb in enumerate(embeddings):
    print(f"Embedding {i+1}:")
    print(f"  Scope: {emb['embedding_scope']}")
    print(f"  Time range: {emb['start_offset_sec']} - {emb['end_offset_sec']} seconds")
    print(f"  Embedding vector (first 5 values): {emb['embedding'][:5]}")
    print()

from typing import List
from aperturedb.DataModels import VideoDataModel, ClipDataModel, DescriptorDataModel, DescriptorSetDataModel
from aperturedb.CommonLibrary import create_connector, execute_query
from aperturedb.Query import generate_add_query
from aperturedb.Query import RangeType
from aperturedb.Connector import Connector
import json

# Define data models for the association of Video, Video Clips, and Embeddings
class ClipEmbeddingModel(ClipDataModel):
    embedding: DescriptorDataModel

class VideoClipsModel(VideoDataModel):
    title: str
    description: str
    clips: List[ClipEmbeddingModel] = []

def create_video_object_with_clips(URL: str, clips, collection):
    video = VideoClipsModel(url=URL, title="Ecommerce v2.5",
                            description="Ecommerce v2.5 video with clips by Marengo26")
    for clip in clips:
        video.clips.append(ClipEmbeddingModel(
            range_type=RangeType.TIME,
            start=clip['start_offset_sec'],
            stop=clip['end_offset_sec'],
            embedding=DescriptorDataModel(
                vector=clip['embedding'], set=collection)
        ))
    return video

video_url = "https://storage.googleapis.com/ad-demos-datasets/videos/Ecommerce%20v2.5.mp4"
clips = embeddings

# Instantiate an ApertureDB client
aperturedb_client = Connector(
    host="workshop.datasets.gcp.cloud.aperturedata.io",
    user="admin",
    password=ADB_PASSWORD
)

# Create a descriptor set (collection)
collection = DescriptorSetDataModel(
    name="marengo26", dimensions=len(clips[0]['embedding']))
q, blobs, c = generate_add_query(collection)
result, response, blobs = execute_query(query=q, blobs=blobs, client=aperturedb_client)
print(f"Descriptor set creation: {result=}, {response=}")

# Create and insert the video object with clips and embeddings
video = create_video_object_with_clips(video_url, clips, collection)
q, blobs, c = generate_add_query(video)
result, response, blobs = execute_query(query=q, blobs=blobs, client=aperturedb_client)
print(f"Video insertion: {result=}, {response=}")

import struct
from aperturedb.Descriptors import Descriptors
from aperturedb.Query import ObjectType
from aperturedb.NotebookHelpers import display_video_mp4
from IPython.display import display

# Generate a text embedding for our search query
text_embedding = twelvelabs_client.embed.create(
  engine_name="Marengo-retrieval-2.6",
  text="Show me the part which has lot of outfits being displayed",
  text_truncate="none"
)

print("Created a text embedding")
print(f" Engine: {text_embedding.engine_name}")
print(f" Embedding: {text_embedding.text_embedding.float[:5]}...")  # Display first 5 values

# Define the descriptor set we'll search in
descriptorset = "marengo26"

# Find similar descriptors to the text embedding
descriptors = Descriptors(aperturedb_client)
descriptors.find_similar(
  descriptorset,
  text_embedding.text_embedding.float,
  k_neighbors=3,
  distances=True
)

# Find connected clips to the descriptors
clip_descriptors = descriptors.get_connected_entities(ObjectType.CLIP)

print(f"Found {len(clip_descriptors)} relevant clips")

# Show the metadata of the clips and the corresponding video segments
for i, clips in enumerate(clip_descriptors, 1):
    print(f"\nResult {i}:")
    for clip in clips:
        print(f"Clip metadata:")
        print(f"  Start time: {clip.start} seconds")
        print(f"  End time: {clip.stop} seconds")
        print(f"  Video URL: {clip.url}")
        
        # Display the video clip
        print("Displaying video clip:")
        display_video_mp4(clips.get_blob(clip))
        print("\n" + "-"*50 + "\n")