Spaces:
Running
Running
| 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() |