add code files

README.md

@@ -1,11 +1,12 @@
 ---
-title: Browser Backend
-emoji: 🔥
-colorFrom: yellow
-colorTo: blue
-sdk: static
+title: Bioscan Updated Ids
+emoji: 👁
+colorFrom: pink
+colorTo: gray
+sdk: gradio
+sdk_version: 5.5.0
+app_file: app.py
 pinned: false
-license: mit
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,122 @@
+import gradio as gr
+import torch
+import numpy as np
+import h5py
+import faiss
+from PIL import Image
+import io 
+import pickle
+import random
+
+def getRandID():
+    indx = random.randrange(0, 396503)
+    return indx_to_id_dict[indx], indx
+
+def chooseImageIndex(indexType):
+    if (indexType == "FlatIP(default)"):
+        return image_index_IP
+    elif (indexType == "FlatL2"):
+        return image_index_L2
+    elif (indexType == "HNSWFlat"):
+        return image_index_HNSW
+    elif (indexType == "IVFFlat"):
+        return image_index_IVF
+    elif (indexType == "LSH"):
+        return image_index_LSH
+
+def chooseDNAIndex(indexType):
+    if (indexType == "FlatIP(default)"):
+        return dna_index_IP
+    elif (indexType == "FlatL2"):
+        return dna_index_L2
+    elif (indexType == "HNSWFlat"):
+        return dna_index_HNSW
+    elif (indexType == "IVFFlat"):
+        return dna_index_IVF
+    elif (indexType == "LSH"):
+        return dna_index_LSH
+
+
+
+def searchEmbeddings(id, mod1, mod2, indexType):
+    # variable and index initialization
+    dim = 768
+    count = 0
+    num_neighbors = 10
+
+    index = faiss.IndexFlatIP(dim)
+
+    # get index
+    if (mod2 == "Image"):
+        index = chooseImageIndex(indexType)
+    elif (mod2 == "DNA"):
+        index = chooseDNAIndex(indexType)
+    
+
+    # search for query
+    if (mod1 == "Image"):
+        query = id_to_image_emb_dict[id]
+    elif (mod1 == "DNA"):
+        query = id_to_dna_emb_dict[id]
+    query = query.astype(np.float32)
+    D, I = index.search(query, num_neighbors)
+
+    id_list = []
+    i = 1
+    for indx in I[0]:
+        id = indx_to_id_dict[indx]
+        id_list.append(id)
+        
+    return id_list
+
+with gr.Blocks() as demo:
+
+    # for hf: change all file paths, indx_to_id_dict as well
+
+    # load indexes
+    image_index_IP = faiss.read_index("big_image_index_FlatIP.index")
+    image_index_L2 = faiss.read_index("big_image_index_FlatL2.index")
+    image_index_HNSW = faiss.read_index("big_image_index_HNSWFlat.index")
+    image_index_IVF = faiss.read_index("big_image_index_IVFFlat.index")
+    image_index_LSH = faiss.read_index("big_image_index_LSH.index")
+
+    dna_index_IP = faiss.read_index("big_dna_index_FlatIP.index")
+    dna_index_L2 = faiss.read_index("big_dna_index_FlatL2.index")
+    dna_index_HNSW = faiss.read_index("big_dna_index_HNSWFlat.index")
+    dna_index_IVF = faiss.read_index("big_dna_index_IVFFlat.index")
+    dna_index_LSH = faiss.read_index("big_dna_index_LSH.index")
+
+    with open("dataset_processid_list.pickle", "rb") as f:
+        dataset_processid_list = pickle.load(f)
+    with open("processid_to_index.pickle", "rb") as f:
+        processid_to_index = pickle.load(f)
+    with open("big_indx_to_id_dict.pickle", "rb") as f:
+        indx_to_id_dict = pickle.load(f)
+
+    # initialize both possible dicts
+    with open("big_id_to_image_emb_dict.pickle", "rb") as f:
+        id_to_image_emb_dict = pickle.load(f)
+    with open("big_id_to_dna_emb_dict.pickle", "rb") as f:
+        id_to_dna_emb_dict = pickle.load(f)
+
+    with gr.Column():
+        with gr.Row():
+            with gr.Column():
+                rand_id = gr.Textbox(label="Random ID:")
+                rand_id_indx = gr.Textbox(label="Index:")
+                id_btn = gr.Button("Get Random ID")
+            with gr.Column():
+                mod1 = gr.Radio(choices=["DNA", "Image"], label="Search From:")
+                mod2 = gr.Radio(choices=["DNA", "Image"], label="Search To:")
+
+        indexType = gr.Radio(choices=["FlatIP(default)", "FlatL2", "HNSWFlat", "IVFFlat", "LSH"], label="Index:", value="FlatIP(default)")
+        process_id = gr.Textbox(label="ID:", info="Enter a sample ID to search for")
+        process_id_list = gr.Textbox(label="Closest 10 matches:" )
+        search_btn = gr.Button("Search") 
+        id_btn.click(fn=getRandID, inputs=[], outputs=[rand_id, rand_id_indx])
+
+    search_btn.click(fn=searchEmbeddings, inputs=[process_id, mod1, mod2, indexType], 
+                     outputs=[process_id_list])
+    
+
+demo.launch()

prepare_index.py

@@ -0,0 +1,105 @@
+from pathlib import Path
+
+import click
+import faiss
+import h5py
+
+ALL_KEY_TYPES = ["dna", "image"]
+ALL_INDEX_TYPES = ["IndexFlatIP", "IndexFlatL2", "IndexIVFFlat", "IndexHNSWFlat", "IndexLSH"]
+EMBEDDING_SIZE = 768
+
+
+def process(input: Path, output: Path, key_type: str, index_type: str):
+    # load embeddings
+    all_keys = h5py.File(input / "extracted_features_of_all_keys.hdf5", "r", libver="latest")[
+        f"encoded_{key_type}_feature"
+    ][:]
+    seen_test = h5py.File(input / "extracted_features_of_seen_test.hdf5", "r", libver="latest")[
+        f"encoded_{key_type}_feature"
+    ][:]
+    unseen_test = h5py.File(input / "extracted_features_of_unseen_test.hdf5", "r", libver="latest")[
+        f"encoded_{key_type}_feature"
+    ][:]
+    seen_val = h5py.File(input / "extracted_features_of_seen_val.hdf5", "r", libver="latest")[
+        f"encoded_{key_type}_feature"
+    ][:]
+    unseen_val = h5py.File(input / "extracted_features_of_unseen_val.hdf5", "r", libver="latest")[
+        f"encoded_{key_type}_feature"
+    ][:]
+
+    # FlatIP and FlatL2
+    if index_type == "IndexFlatIP":
+        test_index = faiss.IndexFlatIP(EMBEDDING_SIZE)
+    elif index_type == "IndexFlatL2":
+        test_index = faiss.IndexFlatL2(EMBEDDING_SIZE)
+    elif index_type == "IndexIVFFlat":
+        # IVFFlat
+        quantizer = faiss.IndexFlatIP(EMBEDDING_SIZE)
+        test_index = faiss.IndexIVFFlat(quantizer, EMBEDDING_SIZE, 128)
+        test_index.train(all_keys)
+        test_index.train(seen_test)
+        test_index.train(unseen_test)
+        test_index.train(seen_val)
+        test_index.train(unseen_val)
+    elif index_type == "IndexHNSWFlat":
+        # HNSW
+        # 16: connections for each vertex. efSearch: depth of search during search. efConstruction: depth of search during build
+        test_index = faiss.IndexHNSWFlat(EMBEDDING_SIZE, 16)
+        test_index.hnsw.efSearch = 32
+        test_index.hnsw.efConstruction = 64
+    elif index_type == "IndexLSH":
+        # LSH
+        test_index = faiss.IndexLSH(EMBEDDING_SIZE, EMBEDDING_SIZE * 2)
+    else:
+        raise ValueError(f"Index type {index_type} is not supported")
+
+    test_index.add(all_keys)
+    test_index.add(seen_test)
+    test_index.add(unseen_test)
+    test_index.add(seen_val)
+    test_index.add(unseen_val)
+
+    faiss.write_index(test_index, str(output / f"bioscan_5m_{key_type}_{index_type}.index"))
+    print("Saved index to", output / f"bioscan_5m_{key_type}_{index_type}.index")
+
+
+@click.command()
+@click.option(
+    "--input",
+    type=click.Path(path_type=Path),
+    default="bioscan-clip-scripts/extracted_features",
+    help="Path to extracted features",
+)
+@click.option(
+    "--output", type=click.Path(path_type=Path), default="bioscan-clip-scripts/index", help="Path to save the index"
+)
+@click.option(
+    "--key-type", "key_type", type=click.Choice(["all", *ALL_KEY_TYPES]), default="all", help="Type of key to use"
+)
+@click.option(
+    "--index-type",
+    "index_type",
+    type=click.Choice(["all", *ALL_INDEX_TYPES]),
+    default="all",
+    help="Type of index to use",
+)
+def main(input, output, key_type, index_type):
+    output.mkdir(parents=True, exist_ok=True)
+
+    if key_type == "all":
+        key_types = ALL_KEY_TYPES
+    else:
+        key_types = [key_type]
+
+    if index_type == "all":
+        index_types = ALL_INDEX_TYPES
+    else:
+        index_types = [index_type]
+
+    for key_type in key_types:
+        for index_type in index_types:
+            process(input, output, key_type, index_type)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,6 @@
+torch
+numpy~=1.26.0
+faiss-cpu==1.8.0
+h5py~=3.9.0
+gradio==5.7.1
+click==8.1.7