Spaces:
Sleeping
Sleeping
| 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 get_image(image1, image2, dataset_image_mask, processid_to_index, idx): | |
| if (idx < 162834): | |
| image_enc_padded = image1[idx].astype(np.uint8) | |
| elif(idx >= 162834): | |
| image_enc_padded = image2[idx-162834].astype(np.uint8) | |
| enc_length = dataset_image_mask[idx] | |
| image_enc = image_enc_padded[:enc_length] | |
| image = Image.open(io.BytesIO(image_enc)) | |
| return image | |
| def searchEmbeddings(id, mod1, mod2): | |
| # variable and index initialization | |
| original_indx = processid_to_index[id] | |
| dim = 768 | |
| num_neighbors = 10 | |
| # get index | |
| index = faiss.IndexFlatIP(dim) | |
| if (mod2 == "Image"): | |
| index = faiss.read_index("image_index.index") | |
| elif (mod2 == "DNA"): | |
| index = faiss.read_index("dna_index.index") | |
| # search index | |
| 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 = [] | |
| for indx in I[0]: | |
| id = indx_to_id_dict[indx] | |
| id_list.append(id) | |
| # get images | |
| image0 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, original_indx) | |
| image1 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][0]) | |
| image2 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][1]) | |
| image3 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][2]) | |
| image4 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][3]) | |
| image5 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][4]) | |
| image6 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][5]) | |
| image7 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][6]) | |
| image8 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][7]) | |
| image9 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][8]) | |
| image10 = get_image(dataset_image1, dataset_image2, dataset_image_mask, processid_to_index, I[0][9]) | |
| # get taxonomic information | |
| # s0 = getTax(original_indx) | |
| # s1 = getTax(I[0][0]) | |
| # s2 = getTax(I[0][1]) | |
| # s3 = getTax(I[0][2]) | |
| # s4 = getTax(I[0][3]) | |
| # s5 = getTax(I[0][4]) | |
| # s6 = getTax(I[0][5]) | |
| # s7 = getTax(I[0][6]) | |
| # s8 = getTax(I[0][7]) | |
| # s9 = getTax(I[0][8]) | |
| # s10 = getTax(I[0][9]) | |
| return id_list, image0, image1, image2, image3, image4, image5, image6, image7, image8, image9, image10 | |
| #s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10 | |
| def getRandID(): | |
| indx = random.randrange(0, 325667) | |
| return indx_to_id_dict[indx], indx | |
| # def getTax(indx): | |
| # s = species[indx] | |
| # g = genus[indx] | |
| # f = family[indx] | |
| # str = "Species: " + s + "\nGenus: " + g + "\nFamily: " + f | |
| # return str | |
| with gr.Blocks(title="Bioscan-Clip") as demo: | |
| # open general files | |
| with open("dataset_image1.pickle", "rb") as f: | |
| dataset_image1 = pickle.load(f) | |
| with open("dataset_image2.pickle", "rb") as f: | |
| dataset_image2 = pickle.load(f) | |
| with open("dataset_processid_list.pickle", "rb") as f: | |
| dataset_processid_list = pickle.load(f) | |
| with open("dataset_image_mask.pickle", "rb") as f: | |
| dataset_image_mask = pickle.load(f) | |
| with open("processid_to_index.pickle", "rb") as f: | |
| processid_to_index = pickle.load(f) | |
| with open("indx_to_id_dict.pickle", "rb") as f: | |
| indx_to_id_dict = pickle.load(f) | |
| # open image files | |
| with open("id_to_image_emb_dict.pickle", "rb") as f: | |
| id_to_image_emb_dict = pickle.load(f) | |
| # open dna files | |
| with open("id_to_dna_emb_dict.pickle", "rb") as f: | |
| id_to_dna_emb_dict = pickle.load(f) | |
| # open taxonomy files | |
| # with open("family.pickle", "rb") as f: | |
| # family = [item.decode("utf-8") for item in pickle.load(f)] | |
| # with open("genus.pickle", "rb") as f: | |
| # genus= [item.decode("utf-8") for item in pickle.load(f)] | |
| # with open("species.pickle", "rb") as f: | |
| # species = [item.decode("utf-8") for item in pickle.load(f)] | |
| with gr.Column(): | |
| process_id = gr.Textbox(label="ID:", info="Enter a sample ID to search for") | |
| process_id_list = gr.Textbox(label="Closest 10 matches:" ) | |
| mod1 = gr.Radio(choices=["DNA", "Image"], label="Search From:") | |
| mod2 = gr.Radio(choices=["DNA", "Image"], label="Search To:") | |
| search_btn = gr.Button("Search") | |
| with gr.Row(): | |
| with gr.Column(): | |
| image0 = gr.Image(label="Original", height=550) | |
| tax0 = gr.Textbox(label="Taxonomy") | |
| 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.Row(): | |
| with gr.Column(): | |
| image1 = gr.Image(label=1) | |
| tax1 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image2 = gr.Image(label=2) | |
| tax2 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image3 = gr.Image(label=3) | |
| tax3 = gr.Textbox(label="Taxonomy") | |
| with gr.Row(): | |
| with gr.Column(): | |
| image4 = gr.Image(label=4) | |
| tax4 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image5 = gr.Image(label=5) | |
| tax5 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image6 = gr.Image(label=6) | |
| tax6 = gr.Textbox(label="Taxonomy") | |
| with gr.Row(): | |
| with gr.Column(): | |
| image7 = gr.Image(label=7) | |
| tax7 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image8 = gr.Image(label=8) | |
| tax8 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image9 = gr.Image(label=9) | |
| tax9 = gr.Textbox(label="Taxonomy") | |
| with gr.Column(): | |
| image10 = gr.Image(label=10) | |
| tax10 = gr.Textbox(label="Taxonomy") | |
| id_btn.click(fn=getRandID, inputs=[], outputs=[rand_id, rand_id_indx]) | |
| search_btn.click(fn=searchEmbeddings, inputs=[process_id, mod1, mod2], | |
| outputs=[process_id_list, image0, image1, image2, image3, image4, image5, image6, image7, image8, image9, image10]) | |
| #tax0, tax1, tax2, tax3, tax4, tax5, tax6, tax7, tax8, tax9, tax10]) | |
| examples = gr.Examples( | |
| examples=[["ABOTH966-22", "DNA", "DNA"], | |
| ["CRTOB8472-22", "DNA", "Image"], | |
| ["PLOAD050-20", "Image", "DNA"], | |
| ["HELAC26711-21", "Image", "Image"]], | |
| inputs=[process_id, mod1, mod2],) | |
| demo.launch() |