File size: 1,447 Bytes
4a56eda 2dd7937 4a56eda 2dd7937 |
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 |
# 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 |