1 import tensorflow as tf
  2 import annoy
  3 import numpy as np
  4 import os
  5 from PIL import Image
  6 import matplotlib.pyplot as plt
  7 import matplotlib.image as mpimg
  8 
  9 class ImageSimilaritySearch:
 10     def __init__(self, image_dir='./uploads/6266/', model_name='ResNet50', feature_dim=2048, num_trees=10, index_file='annoy_index.ann'):
 11         self.image_dir = image_dir
 12         self.model_name = model_name
 13         self.feature_dim = feature_dim
 14         self.num_trees = num_trees
 15         self.base_model = None
 16         self.index = None
 17         self.index_file = index_file
 18         self.image_files = None
 19 
 20         # 加載預(yù)訓(xùn)練的CNN模型
 21         self.load_base_model()
 22 
 23         # 設(shè)置或加載Annoy索引
 24         if os.path.exists(self.index_file):
 25             self.load_annoy_index()
 26         else:
 27             self.setup_annoy_index()
 28 
 29     def load_base_model(self):
 30         # 加載預(yù)訓(xùn)練的CNN模型
 31         self.base_model = tf.keras.applications.__getattribute__(self.model_name)(
 32             weights='imagenet', include_top=False, pooling='avg')
 33     
 34     def setup_annoy_index(self):
 35         # 設(shè)置Annoy索引
 36         self.index = annoy.AnnoyIndex(self.feature_dim, 'euclidean')
 37         self.image_files = self.get_image_paths(self.image_dir)
 38 
 39         for i, image_file in enumerate(self.image_files):
 40             image_path = image_file
 41             features = self.extract_features(image_path)
 42             self.index.add_item(i, features)
 43         self.index.build(self.num_trees)
 44         # 保存Annoy索引到文件
 45         self.index.save(self.index_file)
 46 
 47     def load_annoy_index(self):
 48         # 從文件加載Annoy索引
 49         self.index = annoy.AnnoyIndex(self.feature_dim, 'euclidean')
 50         self.index.load(self.index_file)
 51         self.image_files = self.get_image_paths(self.image_dir)
 52 
 53     def update_annoy_index_with_images(self, new_image_paths):
 54         # 更新Annoy索引，添加新圖片
 55         for i, image_path in enumerate(new_image_paths):
 56             features = self.extract_features(image_path)
 57             self.index.add_item(len(self.image_files) + i, features)
 58 
 59         # 重新構(gòu)建Annoy索引
 60         self.index.build(self.num_trees)
 61 
 62     def get_image_paths(self, image_dir):
 63         # 獲取文件夾下所有圖片路徑
 64         image_paths = []
 65         for root, dirs, files in os.walk(image_dir):
 66             for file in files:
 67                 if file.endswith('.jpg') or file.endswith('.jpeg') or file.endswith('.png'):
 68                     image_path = os.path.join(root, file)
 69                     image_paths.append(image_path)
 70         return image_paths
 71 
 72     def extract_features(self, image_path):
 73         # 加載圖像，并使用模型提取特征向量
 74         img = tf.keras.preprocessing.image.load_img(image_path, target_size=(224, 224))
 75         img_array = tf.keras.preprocessing.image.img_to_array(img)
 76         img_array = tf.keras.applications.resnet50.preprocess_input(img_array)
 77         img_array = np.expand_dims(img_array, axis=0)
 78         features = self.base_model.predict(img_array)
 79         return features.flatten()
 80 
 81     def search_similar_images(self, query_image_path, top_k=5):
 82         # 搜索相似圖像
 83         query_features = self.extract_features(query_image_path)
 84         similar_indices = self.index.get_nns_by_vector(query_features, top_k)
 85         similar_images = [self.image_files[i] for i in similar_indices]
 86         return similar_images
 87 
 88     def show_similar_images(self, similar_images):
 89         # 顯示相似圖像
 90         num_images = len(similar_images)
 91         fig, axes = plt.subplots(1, num_images, figsize=(15, 5))
 92 
 93         for i, image_path in enumerate(similar_images):
 94             ax = axes[i] if num_images > 1 else axes
 95             img = mpimg.imread(image_path)
 96             ax.imshow(img)
 97             ax.axis('off')
 98 
 99         plt.show()
100 
101 if __name__ == "__main__":
102     # 創(chuàng)建ImageSimilaritySearch對(duì)象
103     similarity_search = ImageSimilaritySearch()
104 
105     # 示例用法
106     query_image_path = './uploads/6266/Bracelet/SL101-HseK-101/1711003989672129.jpg'
107     similar_images = similarity_search.search_similar_images(query_image_path)
108     print("Similar images:", similar_images)
109 
110     # 顯示相似圖像
111     # similarity_search.show_similar_images(similar_images)