add tests and fix bugs

Author: rom1504

Diff for: .gitignore

@@ -11,4 +11,5 @@ image_folder
 cat
 embedding_folder
 index_folder
-indices_paths.json
+indices_paths.json
+.coverage

Diff for: clip_retrieval/__init__.py

@@ -0,0 +1,6 @@
+"""clip retrieval"""
+
+from .clip_back import clip_back
+from .clip_filter import clip_filter
+from .clip_index import clip_index
+from .clip_inference import clip_inference

Diff for: clip_retrieval/clip_index.py

@@ -4,13 +4,21 @@
 import fire
 import os
 from distutils.dir_util import copy_tree
-from multiprocessing import Process
+import logging
+
+
+LOGGER = logging.getLogger(__name__)
 
 
 def quantize(emb_folder, index_folder, index_name, max_index_memory_usage, current_memory_available, nb_cores):
     """calls autofaiss to build an index"""
     try:
+        LOGGER.debug(f"starting index {index_name}")
         if os.path.exists(emb_folder):
+            LOGGER.debug(
+                f"embedding path exist, building index {index_name}"
+                f"using embeddings {emb_folder} ; saving in {index_folder}"
+            )
             build_index(
                 embeddings_path=emb_folder,
                 index_path=index_folder + "/" + index_name + ".index",
@@ -19,8 +27,10 @@ def quantize(emb_folder, index_folder, index_name, max_index_memory_usage, curre
                 current_memory_available=current_memory_available,
                 nb_cores=nb_cores,
             )
+            LOGGER.debug(f"index {index_name} done")
     except Exception as e:  # pylint: disable=broad-except
-        print(e)
+        LOGGER.exception(f"index {index_name} failed")
+        raise e
 
 
 def clip_index(
@@ -34,32 +44,22 @@ def clip_index(
     nb_cores=None,
 ):
     """indexes clip embeddings using autofaiss"""
-    p = Process(
-        target=quantize,
-        args=(
-            embeddings_folder + "/" + image_subfolder,
-            index_folder,
-            "image",
-            max_index_memory_usage,
-            current_memory_available,
-            nb_cores,
-        ),
+    quantize(
+        embeddings_folder + "/" + image_subfolder,
+        index_folder,
+        "image",
+        max_index_memory_usage,
+        current_memory_available,
+        nb_cores,
     )
-    p.start()
-    p.join()
-    p = Process(
-        target=quantize,
-        args=(
-            embeddings_folder + "/" + text_subfolder,
-            index_folder,
-            "text",
-            max_index_memory_usage,
-            current_memory_available,
-            nb_cores,
-        ),
+    quantize(
+        embeddings_folder + "/" + text_subfolder,
+        index_folder,
+        "text",
+        max_index_memory_usage,
+        current_memory_available,
+        nb_cores,
     )
-    p.start()
-    p.join()
     if copy_metadata:
         copy_tree(embeddings_folder + "/metadata", index_folder + "/metadata")
 

Diff for: clip_retrieval/clip_inference.py

@@ -230,7 +230,7 @@ def __write_batch(self):
             img_emb_mat = np.concatenate(self.image_embeddings)
             output_path_img = self.img_emb_folder + "/img_emb_" + str(self.batch_num)
 
-            with self.fs.open(output_path_img, "wb") as f:
+            with self.fs.open(output_path_img + ".npy", "wb") as f:
                 npb = BytesIO()
                 np.save(npb, img_emb_mat)
                 f.write(npb.getbuffer())
@@ -242,7 +242,7 @@ def __write_batch(self):
             text_emb_mat = np.concatenate(self.text_embeddings)
             output_path_text = self.text_emb_folder + "/text_emb_" + str(self.batch_num)
 
-            with self.fs.open(output_path_text, "wb") as f:
+            with self.fs.open(output_path_text + ".npy", "wb") as f:
                 npb = BytesIO()
                 np.save(npb, text_emb_mat)
                 f.write(npb.getbuffer())

Diff for: pytest.ini

@@ -0,0 +1,4 @@
+
+[pytest]
+log_cli = 1
+log_cli_level = DEBUG

Diff for: tests/test_end2end.py

@@ -0,0 +1,71 @@
+from img2dataset import download
+from clip_retrieval import clip_inference
+from clip_retrieval import clip_index
+import os
+import pandas as pd
+import shutil
+
+test_list = [
+    ["first", "https://placekitten.com/400/600"],
+    ["second", "https://placekitten.com/200/300"],
+    ["third", "https://placekitten.com/300/200"],
+    ["fourth", "https://placekitten.com/400/400"],
+    ["fifth", "https://placekitten.com/200/200"],
+    [None, "https://placekitten.com/200/200"],
+]
+
+
+def generate_parquet(output_file):
+    df = pd.DataFrame(test_list, columns=["caption", "url"])
+    df.to_parquet(output_file)
+
+
+def test_end2end():
+    current_folder = os.path.dirname(__file__)
+    test_folder = current_folder + "/" + "test_folder"
+    if os.path.exists(test_folder):
+        shutil.rmtree(test_folder)
+    if not os.path.exists(test_folder):
+        os.mkdir(test_folder)
+    url_list_name = os.path.join(test_folder, "url_list")
+    image_folder_name = os.path.join(test_folder, "images")
+
+    url_list_name += ".parquet"
+    generate_parquet(url_list_name)
+
+    download(
+        url_list_name,
+        image_size=256,
+        output_folder=image_folder_name,
+        thread_count=32,
+        input_format="parquet",
+        output_format="webdataset",
+        url_col="url",
+        caption_col="caption",
+    )
+
+    assert os.path.exists(image_folder_name)
+
+    embeddings_folder = os.path.join(test_folder, "embeddings")
+
+    clip_inference(
+        input_dataset=f"{image_folder_name}/00000.tar",
+        output_folder=embeddings_folder,
+        input_format="webdataset",
+        enable_metadata=True,
+        write_batch_size=100000,
+        batch_size=512,
+        cache_path=None,
+    )
+
+    assert os.path.exists(embeddings_folder)
+
+    index_folder = os.path.join(test_folder, "index")
+
+    os.mkdir(index_folder)
+
+    clip_index(embeddings_folder, index_folder=index_folder)
+
+    assert os.path.exists(index_folder + "/image.index")
+    assert os.path.exists(index_folder + "/text.index")
+