From e34be2dd82d85eeb9e09c79e2933a5d6fa1092e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9D=80=E7=81=AB=E7=9A=84=E5=86=B0=E5=9D=97nya?= Date: Fri, 26 Jun 2026 22:25:52 +0800 Subject: [PATCH 1/2] =?UTF-8?q?=E6=B7=BB=E5=8A=A0=E6=9B=B4=E5=A4=9A?= =?UTF-8?q?=E9=BB=98=E8=AE=A4=E6=90=9C=E7=B4=A2=E7=9A=84=E6=89=A9=E5=B1=95?= =?UTF-8?q?=E5=90=8D?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- image_search_core/config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/image_search_core/config.py b/image_search_core/config.py index 17176d0..524a2ef 100644 --- a/image_search_core/config.py +++ b/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/" \ No newline at end of file From afc116d75e7919b4971a7bef7cf3d45a8a8eb60e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9D=80=E7=81=AB=E7=9A=84=E5=86=B0=E5=9D=97nya?= Date: Fri, 26 Jun 2026 22:26:58 +0800 Subject: [PATCH 2/2] =?UTF-8?q?=E9=80=82=E5=BA=94=E6=96=B0=E7=89=88?= =?UTF-8?q?=E6=9C=AC=20transformers=20=E5=BA=93?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- image_search_core/indexer.py | 2 +- image_search_core/searcher.py | 34 +++++++++++++++++----------------- 2 files changed, 18 insertions(+), 18 deletions(-) diff --git a/image_search_core/indexer.py b/image_search_core/indexer.py index c3bef45..c0ccd62 100644 --- a/image_search_core/indexer.py +++ b/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)) diff --git a/image_search_core/searcher.py b/image_search_core/searcher.py index 98203de..79e243b 100644 --- a/image_search_core/searcher.py +++ b/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