AnalyticDB：案例：搭建以圖搜圖系統

背景資訊

以圖搜圖在生活中有著廣泛的應用，當您想擁有在電視中看到的一件美麗裙子或者帥氣球鞋時，可以拍張照片，然後開啟淘寶上傳照片，就可以快速地找到這個商品。或者，想知道一張電影截圖的出處時，只要將圖片粘貼到搜尋引擎的圖搜框中，就可以找到相關電影的資訊。以圖搜圖還可以通過照片在海量的人物相簿中快速地找到目標。當您在使用搜尋引擎的以圖搜圖功能時，是否覺得這種“黑科技”遙不可及呢？其實通過AnalyticDB PostgreSQL版向量資料庫提供的高效向量檢索功能，您只需要使用SQL就可以輕鬆地搭建一套以圖搜圖系統。

以圖搜圖原理介紹

以圖搜圖又稱為反向圖搜 (Reverse Image Search)，是一種基於內容的映像檢索 (Content-based Image Retrieval) 技術。以圖片作為查詢的對象，以圖搜圖系統會在大量的映像記錄中返回與查詢映像內容最接近的記錄。例如，商品以圖搜圖會返回與查詢圖片中主體物品相同或相似的圖片資訊；人臉以圖搜圖會根據圖片中人臉特徵返回目標人物的圖片記錄。

以圖搜圖應用的核心模組有兩個：

• 特徵提模數塊：負責從映像中提取視覺特徵，從而獲得一個高維的特徵向量，在這個高維特徵空間中越相似的映像距離越近。

• 向量檢索模組：負責在海量的映像特徵向量集中快速地尋找與查詢圖片特徵最接近的前k個記錄，並返回。

以圖搜圖的流程圖如下所示。

映像特徵提取

當前主流的特徵提取演算法主要是使用深度學習模型，例如VGG、ResNet、Transformer等模型作為主幹網路，然後使用不同的方法產生特徵。產生特徵常用的方法有三種：

• 最簡單的方法，直接將分類模型（如VGG模型）分類層的前一層輸出作為映像的特徵，這種演算法在以圖搜圖情境中往往召回率不是很高。

• 第二種方法，將模型中介層的特徵經過特殊的方法池化（如RMAC和GeM）和降維從而得到。

• 第三種方法，將預訓練模型在目標資料集上使用專門設計的損失函數進行遷移訓練，以提取特徵。例如，商品以圖搜圖特徵提模數型通常需要在商品資料集上進行遷移學習，以便能更加準確地提取不同商品的視覺特徵。

您可以選擇適合當前使用情境的方法，提取映像的特徵，產生特徵向量。

向量檢索

向量檢索又稱為最近鄰 (Nearest Neighbor Search，NNS) 檢索，主要負責在海量特徵向量中快速地尋找與查詢向量距離最近的k個記錄。雖然可以通過遍曆的方法，依次計算查詢向量與資料庫中所有向量的距離，然後排序，得到結果，但是這種方法的時間複雜度在大規模資料情境下基本無法滿足要求。

在實際的應用情境中，通常使用近似最近鄰檢索 (Approximate Nearest Neighbor，ANN) 的方法，ANN主要利用向量資料分布的特性以犧牲一定檢索精度為代價，快速地返回可能是查詢目標最近鄰的資料。

常見ANN的方法有三種：

• 基於局部敏感雜湊 (LSH) 的方法。

• 基於乘積量化的方法。

• 基於圖的方法。

使用AnalyticDB PostgreSQL向量資料庫實現以圖搜圖

步驟一：特徵向量提取

本文使用的工具如下：

• 程式設計語言為Python 3.8。

• 深度學習架構為Pytorch。

• 資料集為CIFAR100，包含了100類映像，每類包含600張圖片。

• 用於提取特徵的網路為已經預訓練的SqueezeNet。SqueezeNet網路很輕量，輸出的特徵向量為1000維。

1. 建立Python環境。

   # 建議使用Anaconda建立新的Python環境。
   conda create -n adbpg_env python=3.8
   conda activate adbpg_env
   
   pip install torchvision
   pip install matplotlib
   pip install psycopg2cffi

2. 下載CIFAR100資料集並進行預先處理。

   import torch
   import torchvision
   
   from torchvision.transforms import (
       Compose, 
       Resize, 
       CenterCrop, 
       ToTensor, 
       Normalize
   )
   
   preprocess = Compose([
       Resize(256),
       CenterCrop(224),
       ToTensor(),
       Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
   ])
   
   DATA_DIRECTORY = "/Users/XXX/Desktop/vector/CIFAR"
   datasets = {
       "CIFAR100": torchvision.datasets.CIFAR100(DATA_DIRECTORY, transform=preprocess, download=True)
   }

3. （可選）查看下載的資料集。

   import numpy as np
   import matplotlib.pyplot as plt
   from mpl_toolkits.axes_grid1 import ImageGrid
   
   def show_images_from_full_dataset(dset, num_rows, num_cols, indices):        
       im_arrays = np.take(dset.data, indices, axis=0)
       labels = map(dset.classes.__getitem__, np.take(dset.targets, indices))
   
       fig = plt.figure(figsize=(10, 10))
       grid = ImageGrid(
           fig, 
           111,
           nrows_ncols=(num_rows, num_cols),
           axes_pad=0.3)
       for ax, im_array, label in zip(grid, im_arrays, labels):
           ax.imshow(im_array)
           ax.set_title(label)
           ax.axis("off")
   
   dataset = datasets["CIFAR100"]
   show_images_from_full_dataset(dataset, 4, 8, [i for i in range(0, 32)])

   

4. 使用Squeezenet1_1模型批量產生所有圖片的特徵向量，並儲存在特徵向量檔案中。本文特徵向量檔案路徑為/Users/XXX/Desktop/vector/features/CIFAR100/features。

   # 準備資料。
   BATCH_SIZE = 100
   dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE)
   
   # 下載模型。
   model = torchvision.models.squeezenet1_1(pretrained=True).eval()
   
   # 提取特徵向量，寫入features_file_path。
   features_file_path = "/Users/XXX/Desktop/vector/features/CIFAR100/features"
   feature_file = open(features_file_path, 'w')
   img_id = 0
   for batch_number, batch in enumerate(dataloader):
       with torch.no_grad():
           batch_imgs = batch[0]  # 0: images
           batch_labels = batch[1]  # 1: labels
           vector_values = model(batch_imgs).tolist()
   
           for i in range(len(vector_values)):
               img_label = dataset.classes[batch_labels[i].item()]
               # print(img_label)
               feature_file.write(str(img_id) + "|" + img_label + "|")
               
               vector_value = vector_values[i]
               assert len(vector_value) == 1000
   
               for j in range(len(vector_value)):
                   if j == 0:
                       feature_file.write("{")
                       feature_file.write(str(vector_value[j]) + ",")
                   elif j == len(vector_value) - 1:
                       feature_file.write(str(vector_value[j]))
                       feature_file.write("}")
                   else:
                       feature_file.write(str(vector_value[j]) + ",")
               feature_file.write("\n")
               
               img_id = img_id + 1
           print("finished extract feature vector for batch: ", batch_number)
   feature_file.close()

   單張圖片得到的特徵向量形式如下所示。

   [2.67548513424756,2.186723470687866,2.376999616622925,2.3993351459503174,2.833254337310791,
   4.141584873199463,1.0177937746047974,2.0199387073516846,2.436871512298584,1.465838789939880,
   4,10.196249008178711,3.3932418823242188,6.087968826293945,7.661309242248535,7.66005373001098,
   6,5.481011390686035,7.513026237487795,5.552321434020996,4.685927867889404,5.635070323944092,...]

步驟二：AnalyticDB PostgreSQL向量資料庫的資料匯入與查詢

1. 建表並添加向量索引。本文以使用Python的psycopg2cffi庫串連資料庫為例。

   重要

   如您的資料庫沒有開通向量功能，請提交工單聯絡支援人員開通。

   import os
   import psycopg2cffi
   
   # 注意，你可以參照以下代碼設定臨時環境變數。
   # os.environ["PGHOST"] = "XX.XXX.XX.XXX"
   # os.environ["PGPORT"] = "XXXXX"
   # os.environ["PGDATABASE"] = "adbpg_test"
   # os.environ["PGUSER"] = "adbpg_test"
   # os.environ["PGPASSWORD"] = "adbpg_test"
   
   connection = psycopg2cffi.connect(
       host=os.environ.get("PGHOST", "XX.XXX.XX.XXX"),
       port=os.environ.get("PGPORT", "XXXXX"),
       database=os.environ.get("PGDATABASE", "adbpg_test"),
       user=os.environ.get("PGUSER", "adbpg_test"),
       password=os.environ.get("PGPASSWORD", "adbpg_test")
   )
   
   cursor = connection.cursor()
   
   # 用於建立表的SQL語句。
   create_table_sql = """
   CREATE TABLE IF NOT EXISTS public.image_search (
       id INTEGER NOT NULL,
       class TEXT,
       image_vector REAL[],
       PRIMARY KEY(id)
   ) DISTRIBUTED BY(id);
   """
   
   # 修改向量列的儲存格式為PLAIN。
   alter_vector_storage_sql = """
   ALTER TABLE public.image_search ALTER COLUMN image_vector SET STORAGE PLAIN;
   """
   
   # 用於建立向量索引的SQL語句。
   create_indexes_sql = """
   CREATE INDEX ON public.image_search USING ann (image_vector) WITH (dim = '1000', hnsw_m = '100', pq_enable='0');
   """
   
   # 執行上述SQL語句。
   cursor.execute(create_table_sql)
   cursor.execute(alter_vector_storage_sql)
   cursor.execute(create_indexes_sql)
   connection.commit()

2. 將資料集的圖片特徵向量匯入到表中。

   import io
   
   # 定義一個產生器函數，逐行處理檔案中的資料。
   def process_file(file_path):
       with open(file_path, 'r') as file:
           for line in file:
               yield line
   
   # 匯入資料的SQL。
   copy_command = """
   COPY public.image_search (id, class, image_vector)
   FROM STDIN WITH (DELIMITER '|');
   """
   
   # 圖片特徵向量檔案。
   features_file_path = "/Users/XXX/Desktop/vector/features/CIFAR100/features"
   
   # 執行COPY命令。
   modified_lines = io.StringIO(''.join(list(process_file(features_file_path))))
   cursor.copy_expert(copy_command, modified_lines)
   connection.commit()

3. 選擇特徵向量檔案中的一張圖片對應的向量，進行搜尋。例如，搜尋ID為4999的圖片。

   def query_analyticdb(collection_name, vector_name, query_embedding, top_k=20):
   # 建立查詢SQL，返回與查詢向量最相近的圖片，同時計算與查詢向量的相似性。
       query_sql = f"""
       SELECT id, class, l2_distance({vector_name},Array{query_embedding}::real[]) AS similarity
       FROM {collection_name}
       ORDER BY {vector_name} <-> Array{query_embedding}::real[]
       LIMIT {top_k};
       """
   
   # 執行查詢。
       connection = psycopg2cffi.connect(
           host=os.environ.get("PGHOST", "XX.XXX.XX.XXX"),
           port=os.environ.get("PGPORT", "XXXXX"),
           database=os.environ.get("PGDATABASE", "adbpg_test"),
           user=os.environ.get("PGUSER", "adbpg_test"),
           password=os.environ.get("PGPASSWORD", "adbpg_test")
       )
   
       cursor = connection.cursor()
       cursor.execute(query_sql)
       results = cursor.fetchall()
       
       return results
     
   # 選擇一條資料作為query。
   def select_feature(file_path, expect_id):
       with open(file_path, 'r') as file:
           for line in file:
               datas = line.split('|')
               if datas[0] == str(expect_id):
                   vec = '[' + datas[2][1:-2] + ']'
                   return vec
       raise ValueError(f"沒有對應id= {expect_id}的資料")
   
   file_path = "/Users/xxxx/Desktop/vector/features/CIFAR100/features"
   
   # 選取id為4999的圖片。
   query_vector = select_feature(file_path, 4999)
   # 查看這張圖片。
   # show_images_from_full_dataset(dataset, 1, 1, [4999], figsize=(1, 1))
   # print(query_vector)
   
   # 執行查詢。
   results = query_analyticdb("image_search", "image_vector", query_vector)

   ID為4999的圖片如下所示。

   

4. 將查詢結果對應的圖片顯示出來。

   說明

   AnalyticDB PostgreSQL向量資料庫提供的是向量近似最近鄰檢索功能，即加快查詢的速度。

   # 擷取上一步返回結果中的圖片id。
   indices = []
   for item in results:
       indices.append(item[0])
   print(indices)
   
   # 顯示圖片。
   show_images_from_full_dataset(dataset, 4, 5, indices)

   查詢結果如下圖所示。