Update app.py

app.py

@@ -8,23 +8,11 @@ import io
 import pickle
 import random
 
-def getRandID():
-    indx = random.randrange(0, 325667)
-    return indx_to_id_dict[indx]
-
-def getTax(indx):
-    s = species[indx]
-    g = genus[indx]
-    f = family[indx]
-    str = "Species: " + s + "\nGenus: " + g + "\nFamily: " + f
-    return str
-
 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))
@@ -55,7 +43,7 @@ def searchEmbeddings(id, mod1, mod2):
     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])
@@ -68,7 +56,7 @@ def searchEmbeddings(id, mod1, mod2):
     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])
@@ -84,7 +72,18 @@ def searchEmbeddings(id, mod1, mod2):
         
     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
 
-with gr.Blocks() as demo:
+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)
@@ -108,12 +107,9 @@ with gr.Blocks() as demo:
         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)]
+    family = [item.decode("utf-8") for item in dataset_hdf5_all_key["family"][:]]
+    genus = [item.decode("utf-8") for item in dataset_hdf5_all_key["genus"][:]]
+    species = [item.decode("utf-8") for item in dataset_hdf5_all_key["species"][:]]
 
     with gr.Column():
         process_id = gr.Textbox(label="ID:", info="Enter a sample ID to search for")
@@ -167,7 +163,7 @@ with gr.Blocks() as demo:
                 image10 = gr.Image(label=10)
                 tax10 = gr.Textbox(label="Taxonomy")
 
-    id_btn.click(fn=getRandID, inputs=[], outputs=rand_id)
+    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])