适应新版本 transformers 库 & 添加更多默认搜索的扩展名

image_search_core/config.py

@@ -1,4 +1,4 @@
 CONFIG_FILE = "config.json"
 DEFAULT_INDEX_FILE = "image_features.npz"
-IMAGE_EXTENSIONS = ["*.webm", "*.webp"]
+IMAGE_EXTENSIONS = ["*.webm", "*.webp", "*.jpg", "*.png", "*.gif"]
 DEFAULT_IMAGE_DIR = "./stickers/"

image_search_core/indexer.py

@@ -210,7 +210,7 @@ def _process_changes(self, indexed_data, new_paths, modified_paths, deleted_path
                         continue
 
                     inputs = processor(images=image, return_tensors="pt").to(self.device)
-                    features = model.get_image_features(**inputs).cpu().numpy()[0]
+                    features = model.get_image_features(**inputs).pooler_output.cpu().numpy()[0]
                     final_data["features"].append(features)
                     final_data["paths"].append(path)
                     final_data["hashes"].append(calculate_hash(path))

image_search_core/searcher.py

@@ -17,14 +17,14 @@ def __init__(self, index_file: str = DEFAULT_INDEX_FILE, model_path: str = None)
         self.index_file = index_file
         self.model_loader = ModelLoader(model_path)
         self.device = self.model_loader.device
-        
+
         print(f"正在从 '{self.index_file}' 加载索引数据...")
         data = np.load(self.index_file, allow_pickle=True)
         self.image_features = torch.from_numpy(data['features']).to(self.device)
-        
+
         self.image_paths = [os.path.abspath(p) for p in data['paths']]
         self.path_to_idx = {path: i for i, path in enumerate(self.image_paths)}
-        
+
         self.normalized_image_features = F.normalize(self.image_features, p=2, dim=-1)
         print(f"索引加载完毕，包含 {len(self.image_paths)} 张图片。")
 
@@ -34,16 +34,16 @@ def search(self, query: str, top_k: int = 5, negative_query: str = None, similar
         """
         if top_k <= 0 or not (query or similar_image_path):
             return []
-            
+
         model, processor = self.model_loader.load()
-        
+
         with torch.no_grad():
             # 1. 构建正面查询向量
             positive_vectors = []
             if query and query.strip():
                 print(f"\n正在为查询 “{query}” 提取文本特征...")
                 inputs = processor(text=[query], return_tensors="pt").to(self.device)
-                text_features = model.get_text_features(**inputs)
+                text_features = model.get_text_features(**inputs).pooler_output
                 positive_vectors.append(text_features)
 
             if similar_image_path and similar_image_path.strip():
@@ -61,33 +61,33 @@ def search(self, query: str, top_k: int = 5, negative_query: str = None, similar
             # 将文本和图片向量平均，创建混合查询向量
             combined_positive_features = torch.mean(torch.cat(positive_vectors, dim=0), dim=0, keepdim=True)
             normalized_positive_features = F.normalize(combined_positive_features, p=2, dim=-1)
-            
+
             # 2. 计算正面余弦相似度
             positive_scores = (normalized_positive_features @ self.normalized_image_features.T).squeeze()
             final_scores = positive_scores
-            
+
             # 3. 如果有负面查询，处理它
             if negative_query and negative_query.strip():
                 negative_keywords = [kw.strip() for kw in negative_query.split(',') if kw.strip()]
                 if negative_keywords:
                     print(f"正在为排除项 {negative_keywords} 提取文本特征...")
                     neg_inputs = processor(text=negative_keywords, return_tensors="pt", padding=True).to(self.device)
-                    neg_text_features = model.get_text_features(**neg_inputs)
+                    neg_text_features = model.get_text_features(**neg_inputs).pooler_output
                     avg_neg_features = neg_text_features.mean(dim=0, keepdim=True)
                     normalized_neg_text_features = F.normalize(avg_neg_features, p=2, dim=-1)
                     negative_scores = (normalized_neg_text_features @ self.normalized_image_features.T).squeeze()
                     final_scores = positive_scores - negative_scores
-                
+
         # 4. 获取 Top K 结果（支持偏移量）
         total_results_to_fetch = min(top_k + offset, len(self.image_paths))
         if total_results_to_fetch <= offset:
             return []
-            
+
         top_results = torch.topk(final_scores, k=total_results_to_fetch)
-        
+
         scores = top_results.values.cpu().numpy()
         indices = top_results.indices.cpu().numpy()
-        
+
         # 应用偏移量
         paginated_indices = indices[offset:]
         paginated_scores = scores[offset:]
@@ -105,10 +105,10 @@ def search_by_image(self, image_path: str, top_k: int = 5, negative_query: str =
         query_idx = self.path_to_idx.get(image_path)
         if query_idx is None:
             raise ValueError(f"图片 '{os.path.basename(image_path)}' 不在索引中。")
-        
+
         print(f"\n正在以图片 “{os.path.basename(image_path)}” 为基准查找相似项...")
         query_vector = self.normalized_image_features[query_idx].unsqueeze(0)
-        
+
         # 计算正面相似度
         positive_scores = (query_vector @ self.normalized_image_features.T).squeeze()
         final_scores = positive_scores
@@ -121,7 +121,7 @@ def search_by_image(self, image_path: str, top_k: int = 5, negative_query: str =
                 if negative_keywords:
                     print(f"正在为排除项 {negative_keywords} 提取文本特征...")
                     neg_inputs = processor(text=negative_keywords, return_tensors="pt", padding=True).to(self.device)
-                    neg_text_features = model.get_text_features(**neg_inputs)
+                    neg_text_features = model.get_text_features(**neg_inputs).pooler_output
                     avg_neg_features = neg_text_features.mean(dim=0, keepdim=True)
                     normalized_neg_features = F.normalize(avg_neg_features, p=2, dim=-1)
                     negative_scores = (normalized_neg_features @ self.normalized_image_features.T).squeeze()
@@ -140,7 +140,7 @@ def search_by_image(self, image_path: str, top_k: int = 5, negative_query: str =
         for idx, score in zip(indices, scores):
             if idx != query_idx:
                 all_results.append({"path": self.image_paths[idx], "score": float(score)})
-        
+
         # 应用偏移量和 top_k
         start = offset
         end = offset + top_k