Update generate_embeddings.py

generate_embeddings.py

@@ -1,10 +1,26 @@
+# 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'])
-formats = []
 while  index<tasks_df.shape[0]:
     image = load_image(tasks_df['image_path'][index])
     inputs = processor(images=[image], return_tensors="pt").to(model.device)
@@ -13,6 +29,6 @@ while  index<tasks_df.shape[0]:
         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 % 100==0:
+    if index % save_interval==0:
          embeddings_df.to_pickle('image_embeddings.pickle')
     index+=1