Update app.py

app.py

@@ -8,13 +8,8 @@ import os
 import json
 from datetime import datetime
 
-# 配置设备 - 强制使用CPU以适配Hugging Face Spaces免费环境
-device = "cpu"
-
-# 动态获取工作目录
-script_dir = os.getcwd()
-root_dir = os.path.join(script_dir, 'GroceryStoreDataset')
-print(f"数据集根目录: {root_dir}")
+# 配置设备
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # 加载CLIP模型和处理器
 model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
@@ -25,15 +20,12 @@ index = Index(url="https://skilled-duckling-934-us1-vector.upstash.io",
               token="ABgFMHNraWxsZWQtZHVja2xpbmctOTM0LXVzMWFkbWluWkRWalpqUTFPV010T0daaU5DMDBORGMwTFdFMVkyUXRaV1JrTVRjNU1EWmpOekZo")
 
 
-# 加载数据集函数 - 改进路径处理和错误日志
+# 加载数据集函数（保持不变）
 def load_dataset(file_path, root_dir):
     data = []
     print(f"加载数据集文件: {file_path}")
-    
-    # 检查文件是否存在
     if not os.path.exists(file_path):
         raise FileNotFoundError(f"数据集文件不存在: {file_path}")
-    
     with open(file_path, 'r', encoding='utf-8') as f:
         lines = f.readlines()
         for i, line in enumerate(lines):
@@ -42,53 +34,30 @@ def load_dataset(file_path, root_dir):
                 if len(parts) != 3:
                     print(f"第 {i + 1} 行格式错误: {line}")
                     continue
-                    
                 image_path, fine_grained_label, coarse_grained_label = parts
-                
-                # 确保路径格式正确（使用正斜杠）
-                image_path = image_path.replace('\\', '/')
-                
-                # 构建完整路径（移除多余的'dataset'前缀）
-                if image_path.startswith('dataset/'):
-                    image_path = image_path[8:]  # 移除'dataset/'前缀
-                    
+                image_path = image_path.replace('/', os.sep)
                 full_image_path = os.path.join(root_dir, 'dataset', image_path)
-                
-                # 检查文件是否存在并可读取
-                if os.path.exists(full_image_path) and os.access(full_image_path, os.R_OK):
+                if os.path.exists(full_image_path):
                     data.append((full_image_path, int(fine_grained_label), int(coarse_grained_label)))
                 else:
-                    print(f"警告: 文件不存在或不可读 - {full_image_path}")
-                    
+                    print(f"警告: 文件不存在 - {full_image_path}")
             except Exception as e:
                 print(f"解析第 {i + 1} 行时出错: {line}")
                 print(f"错误详情: {e}")
-                
     print(f"成功加载 {len(data)} 个样本")
     return data
 
 
-# 特征提取和向量插入函数
+# 特征提取和向量插入函数（保持不变）
 def insert_images_to_index(data):
     print(f"开始向向量数据库插入 {len(data)} 个图像特征...")
     success_count = 0
     error_count = 0
-    
     for image_path, fine_label, coarse_label in data:
         try:
-            # 验证图像文件存在
-            if not os.path.exists(image_path):
-                print(f"错误: 图像文件不存在 - {image_path}")
-                error_count += 1
-                continue
-                
             image = Image.open(image_path)
             features = extract_image_features(image)
-            
-            # 使用规范化的文件路径作为ID的一部分
-            file_id = os.path.basename(image_path).replace('.', '_')
-            vector_id = f"img_{file_id}_{fine_label}"
-            
+            vector_id = f"img_{os.path.basename(image_path)}_{fine_label}"
             vector = Vector(
                 id=vector_id,
                 vector=features,
@@ -98,14 +67,11 @@ def insert_images_to_index(data):
                     "coarse_label": coarse_label
                 }
             )
-            
             index.upsert(vectors=[vector])
             success_count += 1
-            
         except Exception as e:
             print(f"处理图像 {image_path} 时出错: {e}")
             error_count += 1
-            
     print(f"向量插入完成: 成功 {success_count}, 失败 {error_count}")
 
 
@@ -113,226 +79,161 @@ def extract_image_features(image):
     try:
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
-            
         inputs = processor(images=image, return_tensors="pt").to(device)
-        
         with torch.no_grad():
             image_features = model.get_image_features(**inputs)
-            
         image_features = image_features / image_features.norm(dim=-1, keepdim=True)
         return image_features.cpu().numpy().flatten().tolist()
-        
     except Exception as e:
         print(f"特征提取错误: {e}")
         return [0.0] * 512
 
 
-# 搜索函数 - 改进图像加载和错误处理
+# 搜索函数（保持不变）
 def text_search(query_text, top_k=9, min_similarity=0.0):
     try:
         if not query_text.strip():
             return [(Image.new("RGB", (400, 200), "white"), "请输入搜索文字")]
-            
         text_inputs = processor(text=query_text, return_tensors="pt", padding=True).to(device)
-        
         with torch.no_grad():
             text_features = model.get_text_features(**text_inputs)
             text_features = text_features / text_features.norm(dim=-1, keepdim=True)
-            
         results = index.query(
             vector=text_features.cpu().numpy().flatten().tolist(),
             top_k=top_k,
             include_vectors=True,
             include_metadata=True
         )
-        
         filtered_results = [item for item in results if item.score >= min_similarity]
-        
         if not filtered_results:
             return [(Image.new("RGB", (400, 200), "white"), "无匹配结果")]
-            
         gallery_items = []
-        
         for item in filtered_results[:top_k]:
             metadata = item.metadata
             image_path = metadata["image_path"]
-            
-            # 打印路径用于调试
-            print(f"搜索结果图像路径: {image_path}")
-            
             try:
-                # 验证路径有效性
-                if not image_path or not os.path.exists(image_path):
-                    raise FileNotFoundError(f"路径不存在: {image_path}")
-                    
                 img = Image.open(image_path).convert("RGB")
-                
-            except FileNotFoundError as e:
-                print(f"错误: 找不到图像 - {image_path}")
+            except:
                 img = Image.new("RGB", (200, 200), "white")
-                
-            except Exception as e:
-                print(f"加载图像失败: {image_path}, 错误: {e}")
-                img = Image.new("RGB", (200, 200), "white")
-                
             caption = f"相似度: {item.score:.4f}"
             gallery_items.append((img, caption))
-            
         return gallery_items
-        
     except Exception as e:
         print(f"文字搜索错误: {e}")
         return [(Image.new("RGB", (400, 200), "white"), f"错误: {str(e)}")]
 
 
-# 图像搜索函数 - 改进错误处理
+
+# 图像搜索函数
 def image_search(query_image, top_k=9, min_similarity=0.0):
     try:
         if query_image is None:
             return [(Image.new("RGB", (400, 200), "white"), "请上传搜索图像")]
-            
+
         # 提取图像特征
         image_features = extract_image_features(query_image)
-        
-        # 确保特征向量格式正确
-        if not isinstance(image_features, list):
-            image_features = image_features.tolist()
-            
+
+        # 将列表转换为 PyTorch 张量
+        image_features = torch.tensor(image_features)
+
+        # 归一化处理
+        image_features = image_features / image_features.norm(dim=-1, keepdim=True)
+
         # 使用正确的特征向量进行查询
         results = index.query(
-            vector=image_features,
+            vector=image_features.cpu().numpy().flatten().tolist(),
             top_k=top_k,
             include_vectors=True,
             include_metadata=True
         )
-        
+
         filtered_results = []
-        
         for item in results:
             metadata = item.metadata
             image_path = metadata["image_path"]
-            
+
             # 相似度过滤
             if item.score < min_similarity:
                 continue
-                
+
             filtered_results.append(item)
-            
+
         # 处理空结果
         if not filtered_results:
             return [(Image.new("RGB", (400, 200), "white"), "无匹配结果")]
-            
+
         # 构建Gallery所需的元组列表
         gallery_items = []
-        
         for item in filtered_results[:top_k]:
             metadata = item.metadata
             image_path = metadata["image_path"]
-            
-            # 打印路径用于调试
-            print(f"图像搜索结果路径: {image_path}")
-            
-            try:
-                # 验证路径有效性
-                if not image_path or not os.path.exists(image_path):
-                    raise FileNotFoundError(f"路径不存在: {image_path}")
-                    
-                img = Image.open(image_path).convert("RGB")
-                
-            except FileNotFoundError as e:
-                print(f"错误: 找不到图像 - {image_path}")
-                img = Image.new("RGB", (200, 200), "white")
-                
-            except Exception as e:
-                print(f"加载图像失败: {image_path}, 错误: {e}")
-                img = Image.new("RGB", (200, 200), "white")
-                
-            # 组合分数和标签作为标题
-            caption = f"相似度: {item.score:.4f}"
-            gallery_items.append((img, caption))
-            
+            if image_path:
+                try:
+                    img = Image.open(image_path).convert("RGB")
+                except Exception as e:
+                    print(f"加载图片失败: {image_path}, 错误: {e}")
+                    img = Image.new("RGB", (200, 200), "white")
+
+                # 组合分数和标签作为标题
+                caption = f"相似度: {item.score:.4f}"
+                gallery_items.append((img, caption))
+
         return gallery_items
-        
+
     except Exception as e:
         print(f"图像搜索错误: {e}")
         return [(Image.new("RGB", (400, 200), "red"), f"错误: {str(e)}")]
 
 
-# 初始化向量数据库 - 改进路径验证
+# 初始化向量数据库（保持不变）
 def initialize_vector_db():
+    script_dir = os.path.dirname(os.path.abspath(__file__))
+    root_dir = os.path.join(script_dir, 'GroceryStoreDataset')
     flag_file = os.path.join(root_dir, 'dataset', '.vectors_inserted')
-    
-    # 检查标志文件
     if os.path.exists(flag_file):
         print("发现标志文件，跳过向量数据库检查")
         return
-        
     try:
-        # 测试向量数据库连接
         results = index.query(vector=[0.0] * 512, top_k=1, include_metadata=False)
-        
-        if results and len(results) > 0:
+        if results and len(results.get("results", [])) > 0:
             print("向量数据库已有数据，跳过插入")
             os.makedirs(os.path.dirname(flag_file), exist_ok=True)
-            
             with open(flag_file, 'w') as f:
                 f.write("Vectors already exist")
-                
             return
-            
-        # 验证数据集文件
         train_file = os.path.join(root_dir, 'dataset', 'train.txt')
         val_file = os.path.join(root_dir, 'dataset', 'val.txt')
         test_file = os.path.join(root_dir, 'dataset', 'test.txt')
-        
         for file_path in [train_file, val_file, test_file]:
             if not os.path.exists(file_path):
                 print(f"警告: 数据集文件不存在 - {file_path}")
                 return
-                
-        # 加载数据集
         train_data = load_dataset(train_file, root_dir)
         val_data = load_dataset(val_file, root_dir)
         test_data = load_dataset(test_file, root_dir)
-        
-        # 插入向量
         insert_images_to_index(train_data + val_data + test_data)
-        
-        # 创建标志文件
         os.makedirs(os.path.dirname(flag_file), exist_ok=True)
-        
         with open(flag_file, 'w') as f:
             f.write("Vectors inserted successfully")
-            
     except Exception as e:
         print(f"查询向量数据库失败: {e}")
-        
         if os.path.exists(flag_file):
             print("但发现标志文件，推测数据已插入，跳过插入")
             return
-            
         print("没有标志文件，尝试加载数据并插入（有重复风险）")
-        
-        # 尝试恢复数据加载
-        if 'train_data' not in locals():
+        if train_data is None:
             train_file = os.path.join(root_dir, 'dataset', 'train.txt')
             val_file = os.path.join(root_dir, 'dataset', 'val.txt')
             test_file = os.path.join(root_dir, 'dataset', 'test.txt')
-            
             for file_path in [train_file, val_file, test_file]:
                 if not os.path.exists(file_path):
                     print(f"警告: 数据集文件不存在 - {file_path}")
                     return
-                    
             train_data = load_dataset(train_file, root_dir)
             val_data = load_dataset(val_file, root_dir)
             test_data = load_dataset(test_file, root_dir)
-            
         insert_images_to_index(train_data + val_data + test_data)
-        
-        # 创建标志文件
         os.makedirs(os.path.dirname(flag_file), exist_ok=True)
-        
         with open(flag_file, 'w') as f:
             f.write("Vectors inserted with error handling")
 
@@ -340,11 +241,11 @@ def initialize_vector_db():
 # 主应用界面
 def create_app():
     initialize_vector_db()
-    
+
     with gr.Blocks(title="CLIP图像搜索系统", theme=gr.themes.Soft()) as app:
         gr.Markdown("# CLIP图像搜索系统")
         gr.Markdown("使用文字或图像搜索相似的商品图片")
-        
+
         with gr.Tabs():
             # 文字搜索标签页
             with gr.Tab("文字搜索"):
@@ -355,14 +256,13 @@ def create_app():
                             placeholder="点击下方标签自动填充",
                             interactive=True
                         )
-                        
-                        # 可选标签
+
+                        # 可选标签（使用HTML按钮实现可取消选择）
                         gr.Markdown("### 可选标签")
                         with gr.Row():
                             # 示例标签，可根据实际数据扩展
                             labels = ["apple", "banana", "orange", "vegetables", "fruit"]
                             label_btns = []
-                            
                             for label in labels:
                                 btn = gr.Button(
                                     label,
@@ -370,55 +270,52 @@ def create_app():
                                     elem_classes="tag-btn"
                                 )
                                 label_btns.append(btn)
-                                
                                 # 点击标签时触发的函数
                                 btn.click(
-                                    fn=lambda txt, lbl: lbl if txt != lbl else "",
+                                    fn=lambda txt, lbl: lbl if txt != lbl else "",  # 点击已选标签则清空
                                     inputs=[text_query, gr.Textbox(value=label, visible=False)],
                                     outputs=text_query
                                 )
-                        
+
+
                         # 控制区
                         with gr.Group():
                             gr.Markdown("### 搜索参数")
                             text_top_k = gr.Slider(minimum=1, maximum=21, step=1, value=9, label="最多显示图片数")
                             text_min_sim = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, value=0.0,
                                                      label="最低相似度阈值")
-                        
+
                         text_search_btn = gr.Button("搜索", variant="primary")
-                
+
                 text_output_images = gr.Gallery(label="搜索结果", show_label=True, columns=3, rows=7)
-            
-            # 图像搜索标签页
+
+            # 图像搜索标签页（保持不变）
             with gr.Tab("图像搜索"):
                 with gr.Row():
                     with gr.Column(scale=2):
                         image_query = gr.Image(label="上传搜索图像", type="pil")
-                        
                         with gr.Group():
                             gr.Markdown("### 搜索参数")
                             image_top_k = gr.Slider(minimum=1, maximum=21, step=1, value=9, label="最多显示图片数")
                             image_min_sim = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, value=0.0,
                                                       label="最低相似度阈值")
-                        
                         image_search_btn = gr.Button("搜索", variant="primary")
-                
                 image_output_images = gr.Gallery(label="搜索结果", show_label=True, columns=3, rows=7)
-        
+
         # 文字搜索按钮事件绑定
         text_search_btn.click(
             fn=text_search,
             inputs=[text_query, text_top_k, text_min_sim],
             outputs=text_output_images
         )
-        
+
         # 图像搜索按钮事件绑定
         image_search_btn.click(
             fn=image_search,
             inputs=[image_query, image_top_k, image_min_sim],
             outputs=image_output_images
         )
-        
+
         # 全局样式：标签按钮样式
         gr.Markdown("""
         <style>
@@ -439,7 +336,7 @@ def create_app():
             }
         </style>
         """)
-    
+
     return app