Update app.py

app.py

@@ -7,9 +7,13 @@ from PIL import Image
 import io 
 import pickle
 
-def get_image(file, dataset_image_mask, processid_to_index, idx):
+def get_image(image1, image2, dataset_image_mask, processid_to_index, idx):
     # idx = processid_to_index[query_id]
-    image_enc_padded = file["image"][idx].astype(np.uint8)
+    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))
@@ -47,27 +51,45 @@ def searchEmbeddings(id):
         id_list.append(id)
 
     # get image data
-    dataset_hdf5_all_key = h5py.File('full5m/BIOSCAN_5M.hdf5', "r", libver="latest")['all_keys']
-    dataset_processid_list = [item.decode("utf-8") for item in dataset_hdf5_all_key["processid"][:]]
-    dataset_image_mask = dataset_hdf5_all_key["image_mask"][:]
-    processid_to_index = {pid: idx for idx, pid in enumerate(dataset_processid_list)}
-
-    image1 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][0])
-    image2 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][1])
-    image3 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][2])
-    image4 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][3])
-    image5 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][4])
-    image6 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][5])
-    image7 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][6])
-    image8 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][7])
-    image9 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][8])
-    image10 = get_image(dataset_hdf5_all_key, dataset_image_mask, processid_to_index, I[0][9])
+    # 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("dataset_image1.pickle", "rb") as f:
+    #     dataset_image1 = pickle.load(f)
+    # with open("dataset_image2.pickle", "rb") as f:
+    #     dataset_image2 = pickle.load(f)
+
+    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])
 
     # return id_list, id_list[0], id_list[1], id_list[2], id_list[3], id_list[4], id_list[5], id_list[6], id_list[7], id_list[8], id_list[9], image1, image2, image3, image4, image5, image6, image7, image8, image9, image10
     # return id_list, indx_to_id_dict[I[0][0]], indx_to_id_dict[I[0][1]], indx_to_id_dict[I[0][2]], indx_to_id_dict[I[0][3]], indx_to_id_dict[I[0][4]], indx_to_id_dict[I[0][5]], indx_to_id_dict[I[0][6]], indx_to_id_dict[I[0][7]], indx_to_id_dict[I[0][8]], indx_to_id_dict[I[0][9]]
     return id_list, image1, image2, image3, image4, image5, image6, image7, image8, image9, image10
 
 with gr.Blocks() as demo:
+
+    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("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 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:" )
@@ -90,4 +112,4 @@ with gr.Blocks() as demo:
                      outputs=[process_id_list, image1, image2, image3, image4, image5, image6, image7, image8, image9, image10])
 
 # ARONZ671-20
-demo.launch(share=True)
+demo.launch()