Update app.py

app.py

@@ -8,8 +8,11 @@ import io
 import pickle
 import random
 
-def get_image(file, dataset_image_mask, processid_to_index, idx):
-    image_enc_padded = file[idx].astype(np.uint8)
+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))
@@ -42,17 +45,17 @@ def searchEmbeddings(id, mod1, mod2):
         id_list.append(id)
     
     # get images
-    image0 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, original_indx)
-    image1 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][0])
-    image2 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][1])
-    image3 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][2])
-    image4 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][3])
-    image5 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][4])
-    image6 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][5])
-    image7 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][6])
-    image8 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][7])
-    image9 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][8])
-    image10 = get_image(dataset_hdf5_all_key['image'], dataset_image_mask, processid_to_index, I[0][9])
+    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)
@@ -82,7 +85,10 @@ def getTax(indx):
 
 with gr.Blocks(title="Bioscan-Clip") as demo:
     # open general files
-    dataset_hdf5_all_key = h5py.File('full5m/BIOSCAN_5M.hdf5', "r", libver="latest")['all_keys']
+    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: 
@@ -101,10 +107,12 @@ with gr.Blocks(title="Bioscan-Clip") as demo:
         id_to_dna_emb_dict = pickle.load(f)
 
     # open taxonomy files
-    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 open("tax/family.pickle", "rb") as f:
+        family = [item.decode("utf-8") for item in pickle.load(f)]
+    with open("tax/genus.pickle", "rb") as f:
+        genus= [item.decode("utf-8") for item in pickle.load(f)]
+    with open("tax/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:" )