# Given a DataFrame tasks_df with 'image_path' col that contains all images paths this script will produce a 'image_embeddings.pickle'
# file that contains all the embeddings. You can stop and resume whenever you want, it will restart from the last saved image file
#
import pandas as pd
import sys
from transformers import pipeline
import torch
from transformers import AutoModel, AutoProcessor
from transformers.image_utils import load_image

ckpt = "google/siglip2-so400m-patch16-512"
model = AutoModel.from_pretrained(ckpt, device_map="auto").eval()
processor = AutoProcessor.from_pretrained(ckpt)

tasks_df = # load DataFrame with 'image_path' col that contains all images paths
save_interval = 100 # save embeddings file every save_interval images

try:
    embeddings_df = pd.read_pickle('image_embeddings.pickle')
    index = embeddings_df.shape[0]
except:
    index=0
    embeddings_df = pd.DataFrame(columns=['image_embedding'])
while  index<tasks_df.shape[0]:
    image = load_image(tasks_df['image_path'][index])
    inputs = processor(images=[image], return_tensors="pt").to(model.device)
    
    with torch.no_grad():
        image_embeddings = model.get_image_features(**inputs)    
    new_row = {'image_embedding': image_embeddings}
    embeddings_df = pd.concat([embeddings_df, pd.DataFrame([new_row])], ignore_index=True)
    if index % save_interval==0:
         embeddings_df.to_pickle('image_embeddings.pickle')
    index+=1