File size: 2,495 Bytes
d50ce1c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
from sentence_transformers import SentenceTransformer
from transformers import AutoImageProcessor, AutoModel
from tqdm import tqdm
from PIL import Image
from typing import Iterator, TextIO, List, Dict, Any, Optional, Sequence, Union
import numpy as np
import uuid
from utils import load_json_file
def embed_texts(text_ls:List[str], text_model=None):
if text_model is None:
text_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
text_embeddings = []
for i, text in enumerate(tqdm(text_ls, desc="Embedding text")):
embeds = text_model.encode(text)
text_embeddings.append(embeds)
return np.array(text_embeddings)
def embed_images(image_path_ls:List[str], vision_model=None, vision_model_processor=None):
if vision_model is None or vision_model_processor is None:
vision_model_processor = AutoImageProcessor.from_pretrained('facebook/dinov2-small')
vision_model = AutoModel.from_pretrained('facebook/dinov2-small')
image_embeds_ls = []
for i, frame in enumerate(tqdm(image_path_ls, desc="Embedding image")):
frame = Image.open(frame)
# TODO: add device here
inputs = vision_model_processor(images=frame, return_tensors="pt")
outputs = vision_model(**inputs)
image_embeds_ls.append(outputs.pooler_output)
return np.array([elem.squeeze().detach().numpy() for elem in image_embeds_ls])
def indexing(index, model_stack, vid_metadata_path):
text_model, vision_model, vision_model_processor, _, _ = model_stack
# read metadata file
vid_metadata = load_json_file(vid_metadata_path)
# embed transcripts
vid_trans = [frame['transcript'] for frame in vid_metadata]
transcript_embeddings = embed_texts(text_ls=vid_trans, text_model=text_model)
# embed caption
vid_captions = [frame['caption'] for frame in vid_metadata]
caption_embeddings = embed_texts(text_ls=vid_captions, text_model=text_model)
# embed frames
vid_img_paths = [vid['extracted_frame_path'] for vid in vid_metadata]
frame_embeddings = embed_images(vid_img_paths, vision_model, vision_model_processor)
for ls in [transcript_embeddings, caption_embeddings, frame_embeddings]:
# Prepare metadata
vectors = [
(str(uuid.uuid4()), emb.tolist(), meta) # Generate unique IDs
for emb, meta in zip(ls, vid_metadata)
]
# Upsert vectors into Pinecone
index.upsert(vectors)
|