多模态RAG应用之实现文本检索视频内容

from PIL import Image
import torch
from transformers import AutoProcessor, BridgeTowerProcessor, BridgeTowerModel,BridgeTowerForContrastiveLearning
import numpy as np
from numpy.linalg import norm

def bt_embedding_from_local_pretrained(prompt,image_path):
    #使用的是这个模型：https://huggingface.co/BridgeTower/bridgetower-large-itm-mlm-itc
    model_name = "BridgeTower/bridgetower-large-itm-mlm-itc"
    processor = AutoProcessor.from_pretrained(model_name)
    model = BridgeTowerForContrastiveLearning.from_pretrained(model_name)
    text_only = False

    # 进行预处理 
    if image_path!=Noneand image_path != '':
        image= Image.open(image_path)
        inputs = processor(images=image, text=prompt, return_tensors="pt")
    else: 
        image= Image.open("dummy.jpg")
        inputs = processor(images=image,text=prompt, return_tensors="pt")
        text_only = True

    # 获取模型输出
    with torch.no_grad():
        outputs = model(**inputs)

    # 这里我们假设需要的是池化后的特征
    if text_only:
        embeddings = outputs.text_embeds
    else:
        embeddings = outputs.cross_embeds
    # 转换为列表形式
    embeddings_list = embeddings.squeeze().tolist()

    return embeddings_list

#图片信息
url1='http://farm3.staticflickr.com/2519/4126738647_cc436c111b_z.jpg'
cap1='A motorcycle sits parked across from a herd of livestock'

url2='http://farm3.staticflickr.com/2046/2003879022_1b4b466d1d_z.jpg'
cap2='Motorcycle on platform to be worked on in garage'

url3='http://images.cocodataset.org/val2017/000000360943.jpg'
cap3='a cat laying down stretched out near a laptop'

#这里是上图1 2 3信息
img1 = {
'flickr_url': url1,
'caption': cap1,
'image_path' : 'd:/下载/rag/materials/motorcycle_1.jpg'
}
img2 = {
    'flickr_url': url2,
    'caption': cap2,
    'image_path' : 'd:/下载/rag/materials/motorcycle_2.jpg'
}
img3 = {
    'flickr_url' : url3,
    'caption': cap3,
    'image_path' : 'd:/下载/rag/materials/cat_1.jpg'
}

imgs = [img1, img2, img3]
embeddings = []
for img in [img1, img2, img3]:
    img_path = img['image_path']
    caption = img['caption']
    #这里嵌入图片和文本
    embedding = bt_embedding_from_local_pretrained(caption, img_path) 
    embeddings.append(embedding)

print(embedding) 
print("length:",len(embeddings[0])) 

# 计算余弦计算相似度
def cosine_similarity(vec1, vec2):
    similarity = np.dot(vec1,vec2)/(norm(vec1)*norm(vec2))
    return similarity

ex1_embed = np.array(embeddings[0])
ex2_embed = np.array(embeddings[1])
ex3_embed = np.array(embeddings[2])

# 图片1&2 1&3一组
sim_ex1_ex2 = cosine_similarity(ex1_embed, ex2_embed)
sim_ex1_ex3 = cosine_similarity(ex1_embed, ex3_embed)

print(f"图片1&2的余弦相似度是:{sim_ex1_ex2}")
print(f"图片1&3的余弦相似度是:{sim_ex1_ex3}")

WEBVTT

00:00:00.120 --> 00:00:01.360
today we are flying

00:00:01.360 --> 00:00:04.080
to China with tensions between China and the West,

00:00:04.080 --> 00:00:06.840
the lengthy visa process, and the recent pandemic.

00:00:06.840 --> 00:00:09.840
China has become unknown to the outside world,

00:00:09.840 --> 00:00:11.200
and over the next 4 weeks
......more

#根据字幕提取帧报错为json
def extract_and_save_frames_and_metadata(
        path_to_video, 
        path_to_transcript, 
        path_to_save_extracted_frames,
        path_to_save_metadatas):

    ifnot os.path.exists(path_to_transcript):
        print("字幕不存在") 
        exit()
    
    # metadatas will store the metadata of all extracted frames
    metadatas = []

    # load video using cv2
    video = cv2.VideoCapture(path_to_video)
    # load transcript using webvtt
    trans = webvtt.read(path_to_transcript)
    
    # iterate transcript file
    # for each video segment specified in the transcript file
    for idx, transcript in enumerate(trans):
        # get the start time and end time in seconds
        start_time_ms = str2time(transcript.start)
        end_time_ms = str2time(transcript.end)
        # get the time in ms exactly 
        # in the middle of start time and end time
        # read frame on mid of capthion
        mid_time_ms = (end_time_ms + start_time_ms) / 2
        # get the transcript, remove the next-line symbol
        text = transcript.text.replace("\n", ' ')
        # get frame at the middle time
        video.set(cv2.CAP_PROP_POS_MSEC, mid_time_ms)
        success, frame = video.read()
        if success:
            # if the frame is extracted successfully, resize it
            image = maintain_aspect_ratio_resize(frame, height=350)
            # save frame as JPEG file
            img_fname = f'frame_{idx}.jpg'
            img_fpath = osp.join(
                path_to_save_extracted_frames, img_fname
            )
            imr= cv2.imwrite(img_fpath, image)

            # prepare the metadata
            metadata = {
                'extracted_frame_path': img_fpath,
                'transcript': text,
                'video_segment_id': idx,
                'video_path': path_to_video,
                'mid_time_ms': mid_time_ms,
            }
            metadatas.append(metadata)

        else:
            print(f"ERROR! Cannot extract frame: idx = {idx}")

    # save metadata of all extracted frames
    fn = osp.join(path_to_save_metadatas, 'metadatas.json')
    with open(fn, 'w', encoding='utf-8') as outfile:
        json.dump(metadatas, outfile, ensure_ascii=False, indent=4)
    return metadatas

#视频存储文件夹
vid1_dir = "./shared_data/videos/travelcn2024"
vid1_filepath = os.path.join(vid1_dir, "travel_china2024.mp4")
vid1_transcript_filepath = os.path.join(vid1_dir, "captions.vtt")

# output paths to save extracted frames and their metadata 
extracted_frames_path = osp.join(vid1_dir, 'extracted_frame')
metadatas_path = vid1_dir

# create these output folders if not existing
Path(extracted_frames_path).mkdir(parents=True, exist_ok=True)
Path(metadatas_path).mkdir(parents=True, exist_ok=True)

## 1、Jenerate frames and metadatas
# call the function to extract frames and metadatas
metadatas = extract_and_save_frames_and_metadata(
                vid1_filepath, 
                vid1_transcript_filepath,
                extracted_frames_path,
                metadatas_path,
            )
print(metadatas)

[
    {
        "extracted_frame_path": "./shared_data/videos/travelcn2024\\extracted_frame\\frame_0.jpg",
        "transcript": "today we are flying",
        "video_segment_id": 0,
        "video_path": "./shared_data/videos/travelcn2024\\travel_china2024.mp4",
        "mid_time_ms": 740.0
    },
    {
        "extracted_frame_path": "./shared_data/videos/travelcn2024\\extracted_frame\\frame_1.jpg",
        "transcript": "to China with tensions between China and the West,",
        "video_segment_id": 1,
        "video_path": "./shared_data/videos/travelcn2024\\travel_china2024.mp4",
        "mid_time_ms": 2720.0
    },
    {
        "extracted_frame_path": "./shared_data/videos/travelcn2024\\extracted_frame\\frame_2.jpg",
        "transcript": "the lengthy visa process, and the recent pandemic.",
        "video_segment_id": 2,
        "video_path": "./shared_data/videos/travelcn2024\\travel_china2024.mp4",
        "mid_time_ms": 5460.0
    }
    ......
]

# 处理元数据信息
def preprocess_transcript(vid_metadata_path):    
    if not os.path.exists(vid_metadata_path):
        return

    vid_metadata = load_json_file_GBK(vid_metadata_path)
    vid_trans = [vid['transcript'] for vid in vid_metadata]
    vid_img_path = [vid['extracted_frame_path'] for vid in vid_metadata]

    # 1、原字幕处理成具有上下文信息的新数组，数组长度一样
    n = 7 #这里取7
    updated_vid_trans = [
    ' '.join(vid_trans[i-int(n/2) : i+int(n/2)]) if i-int(n/2) >= 0 else
    ' '.join(vid_trans[0 : i + int(n/2)]) for i in range(len(vid_trans))
    ]

    for i in range(len(updated_vid_trans)):
        vid_metadata[i]['transcript'] = updated_vid_trans[i]

    print(f'第六句处理前:\n"{vid_trans[6]}"')
    print(f'第六句处理后:\n"{updated_vid_trans[6]}"')

    return vid_trans,updated_vid_trans,vid_img_path,vid_metadata

vid_metadata_path=osp.join(metadatas_path, 'metadatas.json')
#这里最终会返回处理后的json元数据，后面需要保存的
vid_trans,updated_vid_trans,vid_img_path,vid_metadata=preprocess_transcript(vid_metadata_path)

rom typing import List
from langchain_core.embeddings import Embeddings
from langchain_core.pydantic_v1 import (
    BaseModel,
)
from utils import encode_image, bt_embedding_from_prediction_guard, bt_embedding_from_local_pretrained
from tqdm import tqdm


class BridgeTowerEmbeddings(BaseModel, Embeddings):
    """ BridgeTower embedding model """
        
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = []
        for text in texts:
            embedding= bt_embedding_from_local_pretrained(text,"")
            embeddings.append(embedding)
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        return self.embed_documents([text])[0]

    def embed_image_text_pairs(self, texts: List[str], images: List[str], batch_size=2) -> List[List[float]]:
        # the length of texts must be equal to the length of images
        assert len(texts)==len(images), "the len of captions should be equal to the len of images"

        embeddings = []
        for path_to_img, text in tqdm(zip(images, texts), total=len(texts)):
            embedding= bt_embedding_from_local_pretrained(text,path_to_img)
            embeddings.append(embedding)
        return embeddings

from typing import Any, Iterable, List, Optional
from langchain_core.embeddings import Embeddings
import uuid
from langchain_community.vectorstores.lancedb import LanceDB

class MultimodalLanceDB(LanceDB):    
    def __init__(
        self,
        connection: Optional[Any] = None,
        embedding: Optional[Embeddings] = None,
        uri: Optional[str] = "/tmp/lancedb",
        vector_key: Optional[str] = "vector",
        id_key: Optional[str] = "id",
        text_key: Optional[str] = "text",
        image_path_key: Optional[str] = "image_path", 
        table_name: Optional[str] = "vectorstore",
        api_key: Optional[str] = None,
        region: Optional[str] = None,
        mode: Optional[str] = "append",
    ):
        super(MultimodalLanceDB, self).__init__(connection, embedding, uri, vector_key, id_key, text_key, table_name, api_key, region, mode)
        self._image_path_key = image_path_key

#插入文本<-> 图片对应记录
    def add_text_image_pairs(
        self,
        texts: Iterable[str],
        image_paths: Iterable[str],
        metadatas: Optional[List[dict]] = None,
        ids: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> List[str]:
        # the length of texts must be equal to the length of images
        assert len(texts)==len(image_paths), "the len of transcripts should be equal to the len of images"

        # Embed texts and create documents
        docs = []
        ids = ids or [str(uuid.uuid4()) for _ in texts]
        embeddings = self._embedding.embed_image_text_pairs(texts=list(texts), images=list(image_paths))  # type: ignore
        for idx, text in enumerate(texts):
            embedding = embeddings[idx]
            metadata = metadatas[idx] if metadatas else {"id": ids[idx]}
            docs.append(
                {
                    self._vector_key: embedding,
                    self._id_key: ids[idx],
                    self._text_key: text,
                    self._image_path_key : image_paths[idx],
                    "metadata": metadata,
                }
            )

        if'mode'in kwargs:
            mode = kwargs['mode']
        else:
            mode = self.mode
        if self._table_name in self._connection.table_names():
            tbl = self._connection.open_table(self._table_name)
            if self.api_key is None:
                tbl.add(docs, mode=mode)
            else:
                tbl.add(docs)
        else:
            self._connection.create_table(self._table_name, data=docs)
        return ids                      

#定义个类函数，供外部调用插入文本->图片嵌入向量数据
    @classmethod
    def from_text_image_pairs(
        cls,
        texts: List[str],
        image_paths: List[str],
        embedding: Embeddings,
        metadatas: Optional[List[dict]] = None,
        connection: Any = None,
        vector_key: Optional[str] = "vector",
        id_key: Optional[str] = "id",
        text_key: Optional[str] = "text",
        image_path_key: Optional[str] = "image_path",
        table_name: Optional[str] = "vectorstore",
        **kwargs: Any,
    ):

        instance = MultimodalLanceDB(
            connection=connection,
            embedding=embedding,
            vector_key=vector_key,
            id_key=id_key,
            text_key=text_key,
            image_path_key=image_path_key,
            table_name=table_name,
        )
        instance.add_text_image_pairs(texts, image_paths, metadatas=metadatas, **kwargs)

        return instance

#初始化数据配置
LANCEDB_HOST_FILE = "./shared_data/.lancedb_test"
TBL_NAME = "test_tbl"
db = lancedb.connect(LANCEDB_HOST_FILE)
#示例化一个表实例
tbl = db.open_table(TBL_NAME)

#先打印看看操作前的表记录数
print(f"Before adding data, There are {tbl.to_pandas().shape[0]} rows in the table")
embedder = BridgeTowerEmbeddings()
dbInstance = MultimodalLanceDB.from_text_image_pairs(
                texts=updated_vid_trans,
                image_paths=vid_img_path,
                embedding=embedder,
                metadatas=vid_metadata,
                connection=db,
                table_name=TBL_NAME,
                mode="append", #overwrite
)
tbl = db.open_table(TBL_NAME)
print(f"After adding data, There are {tbl.to_pandas().shape[0]} rows in the table")

from mm_rag.embeddings.bridgetower_embeddings import (
    BridgeTowerEmbeddings
)
from mm_rag.vectorstores.multimodal_lancedb import MultimodalLanceDB
import lancedb
from PIL import Image
from utils import bt_embedding_from_local_pretrained
import webbrowser
import os

#初始化查询参数
embedder = BridgeTowerEmbeddings()
LANCEDB_HOST_FILE = "./shared_data/.lancedb_test"
TBL_NAME = "test_tbl"
db = lancedb.connect(LANCEDB_HOST_FILE)
#创建
vectorstore = MultimodalLanceDB(
        uri=LANCEDB_HOST_FILE, 
        embedding=embedder, 
        table_name=TBL_NAME
    )

tbl = db.open_table(TBL_NAME)

#搜索图片帧，size是搜索结果的数量
def search_by_prompt(prompt,size): 
    retriever = vectorstore.as_retriever(
        search_type='similarity', 
        search_kwargs={"k": size}
    )
    results = retriever.invoke(prompt)
    print(f"Searching for \"{prompt}\"*********************************************\n")
    #打印出搜索结果
    # display_retrieved_results(results)
    #打开搜索结果的图片
    for i in range(len(results)):
        print(results[i])
        image_path=results[i].metadata['extracted_frame_path']
        Image.open(image_path).show()

#这里要搜索：一个婴儿和大人
search_by_prompt("A toddler and an adult",2)

<!DOCTYPE html>
<head>
    <title>视频片段</title>
</head>
<body>
    <div id="app">
        <video id="video1" :src="video_url" controls autoplay muted height="720" @timeupdate="timeupdate($event)" @play="checkStartTime($event)" > </video>
    </div>
    <script src="https://unpkg.com/vue@3/dist/vue.global.js"></script>
    <script>
        const { createApp } = Vue

        createApp({
            data() {
                return {
                    video_url: "./videos/video2/toddler_in_playground.mp4",
                    start: 2,
                    end: 5
                }
            }
            , methods: {
                timeupdate(event) {
                    const videoElement = event.target;
                    if (videoElement.currentTime > this.end) {
                        videoElement.pause();
                    }
                },
                checkStartTime(event) {
                    const videoElement = event.target;
                    if (videoElement.currentTime < this.start) {
                        videoElement.currentTime = this.start;
                    }
                }
            }, mounted() {
                const urlParams = new URLSearchParams(window.location.search);
                const ms = urlParams.get('ms');
                const vpath = urlParams.get('vpath');
                const vname = urlParams.get('vname');
                this.start = (ms / 1000) - 3
                this.end = ms / 1000 + 3
                if (this.start < 0) {
                    this.start = 0
                }
                console.log(this.start, this.end,vpath,vname);
                this.video_url =`./videos/${vpath}/${vname}`;
            }
        }).mount('#app')

    </script>
</body>
</html

#搜索视频片段，搜到后调用浏览器播放对应视频片段
def search_segment_by_prompt(prompt): 
    retriever = vectorstore.as_retriever(
        search_type='similarity', 
        search_kwargs={"k": 1}
    )
    results = retriever.invoke(prompt)
    ms=results[0].metadata['mid_time_ms']
    video_path=results[0].metadata['video_path']
    print(ms,video_path)
    
    # split the path and get the last two part of them
    normalized_path = os.path.normpath(video_path)
    parts = normalized_path.split(os.sep)
    result = parts[-2:]
    #这里会调起上面的video.html，播放搜索到的视频片段
    url = f'http://127.0.0.1:5500/deeplearing_rag/lesson/shared_data/video.html?ms={ms}&vpath={result[0]}&vname={result[1]}'
    webbrowser.open(url)

#搜一个骑摩托车片段
search_segment_by_prompt("riding motorcycle")