'app.py

app.py

@@ -1,8 +1,8 @@
 import time
 import os
-
-
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
+import zipfile
+from io import BytesIO
 from PIL import Image
 import numpy as np
 import argparse
@@ -40,17 +40,17 @@ def get_parser():
         cfg = config.merge_cfg_from_list(cfg, args.opts)
     return cfg
 args = get_parser()
-data_dir = './LuojiaHOG(demo).json'
-imgs_folder = './image/'
+data_dir = './LuojiaHOG(best)_.json'
+imgs_folder = 'image/'
 
 # image_id = 'sample44_1641.jpg'
 # model_path = './rsvit.pth'
 
-with open('image_features.pkl', 'rb') as f:
+with open('image_features_best.pkl', 'rb') as f:
     image_dict = pickle.load(f)
     image_feat = np.array(list(image_dict.values()))
 f.close()
-with open('text_features.pkl', 'rb') as f:
+with open('text_features_best.pkl', 'rb') as f:
     text_dict = pickle.load(f)
     text_feat = np.array(list(text_dict.values()))
 f.close()
@@ -169,8 +169,20 @@ def search(text_query, image_query, top_k: int = 10):
     #     {"Score": [torch.round(torch.tensor(value)*100)/100 for value in top_k_scores], "Title": top_k_titles, "Text": top_k_texts}
     # )
     # 获取图像
+
+
     if text_query != None:
-        image_output = [Image.open(imgs_folder + img.replace('_','/')) for img in info]
+        # image_output = [Image.open(imgs_folder + img.replace('_','/')) for img in info]
+        image_output = []
+        for img in info:
+            sample_name = img.split('_')[0]
+            image_path = imgs_folder + sample_name + '.zip'
+            with zipfile.ZipFile(image_path, 'r') as zip_ref:
+                # 读取图像文件
+                with zip_ref.open(img.replace('_', '/')) as image_file:
+                    # 将读取的字节流转换为图像
+                    image = Image.open(BytesIO(image_file.read()))
+                    image_output.append(image)
     else:
         image_output = []
 
@@ -282,7 +294,6 @@ def txt_search(text_query, top_k: int = 10):
     # query_embedding = model.encode(query)
 
     text = tokenize(text_query, 328)
-
     query_vector = model.text_encode(text)
     index = text_index