【腾讯云云上实验室】《手把手带你 5 分钟构建以图搜图系统》

curl -i -X POST \
  -H 'Content-Type: application/json' \
  -H 'Authorization: Bearer account=root&api_key=A5VOgsMpGWJhUI0WmUbY********************' \
  http://10.0.X.X:80/database/create \
  -d '{
    "database": "db-test"
}'

{"code":0,"msg":"operation success","affectedCount":1}

mkdir -p image-search
cd image-search

$ python -V
Python 3.9.0

$ python -m venv venv

source venv/bin/activate

.\venv\Scripts\activate

python -m pip install -q towhee opencv-python pillow tcvectordb

curl -L https://github.com/towhee-io/examples/releases/download/data/reverse_image_search.zip -O
unzip -q -o reverse_image_search.zip

class TcvdbClient(PyOperator):

    def __init__(self, host: str, port: str, username: str, key: str, dbName: str, collectionName: str,
                 timeout: int = 20):
        """
        初始化客户端
        """
        # 创建客户端时可以指定 read_consistency，后续调用 sdk 接口的 read_consistency 将延用该值
        self.collectionName = collectionName
        self.db_name = dbName
        self._client = tcvectordb.VectorDBClient(url="http://" + host + ":" + port, username=username, key=key,
                                                 read_consistency=ReadConsistency.EVENTUAL_CONSISTENCY, timeout=timeout)

# tcvdb parameters
HOST = 'lb-xxxx.clb.ap-beijing.tencentclb.com'
PORT = '10000'
DB_NAME = 'image-search'
COLLECTION_NAME = 'reverse_image_search'

PASSWORD = 'xxxx'
USERNAME = 'root'

# path to csv (column_1 indicates image path) OR a pattern of image paths
INSERT_SRC = 'reverse_image_search.csv'

test_vdb = TcvdbClient(host=HOST, port=PORT, key='6qlvBkF0xAgZoN7VJbcwLCqrxoSS4J63Q7mu4RgF', username='root',
                           collectionName=COLLECTION_NAME, dbName=DB_NAME)

class TcvdbClient(PyOperator):
    def create_db_and_collection(self):
        database = self.db_name
        coll_embedding_name = self.collectionName
        coll_alias = self.collectionName + "-alias"

        # 创建 DB
        db = self._client.create_database(database)

        # 构建 Collection
        index = Index()
        index.add(VectorIndex('vector', 2048, IndexType.HNSW, MetricType.COSINE, HNSWParams(m=16, efconstruction=200)))
        index.add(FilterIndex('id', FieldType.String, IndexType.PRIMARY_KEY))
        index.add(FilterIndex('path', FieldType.String, IndexType.FILTER))

        # 创建 Collection
        db.create_collection(
            name=coll_embedding_name,
            shard=3,
            replicas=0,
            description='image embedding collection',
            index=index,
            embedding=None,
            timeout=20
        )

test_vdb.create_db_and_collection()

# Initialize model
# model = models.resnet50(pretrained=True)
model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)

# Set model to eval mode
model = model.eval()
model = torch.nn.Sequential(*(list(model.children())[:-1]))

# Embedding: Function to extract features from an image
def extract_features(image_path):
    transform = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])
    # Read the image Ensure the image is read as RGB
    img = Image.open(image_path).convert('RGB')
    img = transform(img)
    img = img.unsqueeze(0)
    feature = model(img)
    # Reshape the features to 2D
    feature = feature.view(feature.shape[0], -1)
    return feature.detach().numpy()

# path to csv (column_1 indicates image path) OR a pattern of image paths
INSERT_SRC = 'reverse_image_search.csv'

# Insert data：Read the CSV file to get all image paths, process each row and insert the features into the TCVDB
for image_path in load_image(INSERT_SRC):
    # 获取图片的向量
    features = extract_features(image_path)
    # 插入向量数据库
    tcvdb_client.upsert(image_path, features)

class TcvdbClient:
    def upsert_data(self, document_list):
        # 获取 Collection 对象
        db = self._client.database(self.db_name)
        coll = db.collection(self.collectionName)
        # upsert 写入数据，可能会有一定延迟
        # 1. 支持动态 Schema，除了 id、vector 字段必须写入，可以写入其他任意字段；
        # 2. upsert 会执行覆盖写，若文档id已存在，则新数据会直接覆盖原有数据(删除原有数据，再插入新数据)
        coll.upsert(documents=document_list)

    def upsert(self, path, item):
        """
        Insert one row to Tcvdb.

        Args:
        path (`str`):
            The path of the image to insert into Tcvdb.
        item (`np.ndarray`):
            The feature vector of the image.
        """

        # Generate a random UUID and convert to string
        document_list = [
            Document(
                id=str(uuid.uuid4()),
                path=path,
                vector=vector),
        ]

        self.upsert_data(document_list)
        return

class TcvdbClient:
    def search(self, query: 'ndarray'):
        tcvdb_result = self.query_data(query)

        result = []
        # 过滤只显示指定的变量
        for hit in tcvdb_result[0]:
            row = []
            row.extend([hit["path"], hit["score"]])
            result.append(row)
        return result

    def query_data(self, query: []):
        # 获取 Collection 对象
        db = self._client.database(self.db_name)
        coll = db.collection(self.collectionName)
        vector = query
        vector = list(map(float, query.ravel()))

        # search
        # search 提供按照 vector 搜索的能力
        # 其他选项类似 search 接口
        # 批量相似性查询，根据指定的多个向量查找多个 Top K 个相似性结果
        res = coll.search(
            vectors=[vector], # 指定检索向量，最多指定20个
            retrieve_vector=False, # 是否需要返回向量字段，False：不返回，True：返回
            limit=10, # 指定 Top K 的 K 值
        )
        # 输出相似性检索结果，检索结果为二维数组，每一位为一组返回结果，分别对应search时指定的多个向量
        print_object(res)
        return res

# Test Image Path
TEST_IMAGE_PATH = './test/goldfish/*.JPEG'

def search_similar_image(path):
    for query_image in load_image(path):
        features = extract_features(query_image)
        search_res = tcvdb_client.search(features)
        # Process the search results
        # 获取 'pred' 字段的值，类如['/root/image-search/reverse_image_search/train/cuirass/n03146219_11082.JPEG',
        # '/root/image-search/reverse_image_search/train/apiary/n02727426_948.JPEG']
        pred = [str(Path(res[0]).resolve()) for res in search_res]
        print("Query image:", query_image)
        print("Search results:", pred)
        return pred

# Search for example query image(s), process each query image and search in the TCVDB
search_similar_image(TEST_IMAGE_PATH)

Query image: test/Afghan_hound/n02088094_4261.JPEG
Search results: ['/root/image-search/reverse_image_search/train/Afghan_hound/n02088094_6533.JPEG', '/root/image-search/reverse_image_search/train/Afghan_hound/n02088094_2164.JPEG', '/root/image-search/reverse_image_search/train/soft-coated_wheaten_terrier/n02098105_6619.JPEG', '/root/image-search/reverse_image_search/train/Afghan_hound/n02088094_5911.JPEG', '/root/image-search/reverse_image_search/train/soft-coated_wheaten_terrier/n02098105_400.JPEG', '/root/image-search/reverse_image_search/train/Bouvier_des_Flandres/n02106382_9318.JPEG', '/root/image-search/reverse_image_search/train/Afghan_hound/n02088094_5532.JPEG', '/root/image-search/reverse_image_search/train/soft-coated_wheaten_terrier/n02098105_82.JPEG', '/root/image-search/reverse_image_search/train/Afghan_hound/n02088094_6565.JPEG', '/root/image-search/reverse_image_search/train/Afghan_hound/n02088094_14463.JPEG']

import gradio as gr
# webui
iface = gr.Interface(
    fn=search_similar_image,
    inputs=gr.components.Textbox(value='test/goldfish/*.JPEG'),
    outputs=gr.Gallery(label="最终的结果图片").style(height='auto'),
    title='Tencent vector db 案例: 以图搜图',
)
iface.launch()

[
  '/root/image-search/reverse_image_search/train/cuirass/n03146219_11082.JPEG',
  '/root/image-search/reverse_image_search/train/loudspeaker/n03691459_40992.JPEG'
]