feature: add image feature extraction script

embed.py

@@ -0,0 +1,50 @@
+import argparse
+import csv
+import os
+
+from PIL import Image
+
+from utils import load_model
+
+
+def main(args):
+    root = args.image_path
+    files = list(os.listdir(root))
+    for model_name in ["koclip", "koclip/koclip-large"]:
+        counter = 0
+        images = []
+        image_ids = []
+        model, processor = load_model(f"koclip/{model_name}")
+        while counter < len(files):
+            if counter != 0 and counter % args.batch_size == 0:
+                inputs = processor(text=[""], images=images, return_tensors="jax", padding=True)
+                features = model(**inputs).image_embeds
+                with open(os.path.join(args.out_path, f"{model_name}.tsv", "w+")) as f:
+                    writer = csv.writer(f, delimiter="\t")
+                    for image_id, feature in zip(image_ids, features):
+                        writer.writerow([image_id, ",".join(feature)])
+                images = []
+                image_ids = []
+            else:
+                file_ = files[counter]
+                image = Image.open(os.path.join(root, file_))
+                images.append(image)
+                image_ids.append(file_)
+                counter += 1
+        
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--batch_size", default=16)
+    parser.add_argument("--image_path", default="images")
+    parser.add_argument("--out_path", default="features")
+    args = parser.parse_args()
+    main(args)
+
+            
+
+            
+
+
+