move embedder to module

1 files changed, 18 insertions, 0 deletions

diff --git a/src/embedders/imgbeddings-api.py b/src/embedders/imgbeddings-api.py
new file mode 100644
index 0000000..0c890b5
--- /dev/null
+++ b/src/embedders/imgbeddings-api.py
@@ -0,0 +1,18 @@
+from PIL import Image
+from imgbeddings import imgbeddings
+import sys, os
+import json as j
+
+b = imgbeddings()
+
+paths = sys.argv[1:]
+batch_size = 8
+
+for i in range(0, len(paths), batch_size):
+    fs = paths[i:i+batch_size]
+
+    ims = [Image.open(open(f, "rb")) for f in fs]
+    for emb in b.to_embeddings(ims).tolist():
+        print(j.dumps(emb))
+
+sys.stderr.write("\n")