python 使用 imagehash 庫生成 ahash，並存入 milvus

ponponon發表於2022-11-24
Python

有一個需求：計算圖片的相似度

需要解決兩個問題：

  • 生成 ahash
  • 儲存和計算 ahash 之間的距離

生成 ahash

『生成 ahash』 選用 python 下面的一個 imagehash 庫。（github：https://github.com/JohannesBu...）

from io import BytesIO
import numpy
import imagehash
from PIL import Image


def create_vector(file: BytesIO) -> bytes:
    image = Image.open(file)
    hash = imagehash.average_hash(image)

    _vector = []

    for h in hash.hash:
        _vector.extend(h)

    vector = bytes(
        numpy.packbits(
            [
                int(v)
                for v in _vector
            ],
            axis=-1
        ).tolist()
    )

    return vector

create_vector 函式輸出的型別是 bytes，就是二進位制序列

imagehash.average_hash(image) 輸出的 hash 物件，hash 物件有一個 hash 屬性，這個屬性的型別是 list[list[bool]]
列印出來就是長下面這樣子，其實就是一個 8x8=64 bit 的序列

[[False False False False False False False False]
 [ True False False False  True False False False]
 [False False  True  True  True  True False False]
 [False False False  True  True False  True  True]
 [False False  True  True  True False False False]
 [False  True  True  True  True False False False]
 [False  True  True  True  True False  True  True]
 [False False False  True  True False  True  True]]

向量資料庫

『儲存和計算 ahash 之間的距離』選用 milvus

建立集合

定義集合：

import settings
from pymilvus import (
    connections,
    Collection,
    FieldSchema,
    CollectionSchema,
    DataType,
)
from loggers import logger

connections.connect(
    host=settings.MILVUS_CONFIG.host,
    port=settings.MILVUS_CONFIG.port,
)

schema = CollectionSchema([
    FieldSchema("id", DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema("meta_id", DataType.INT64),
    FieldSchema("company_id", DataType.INT64),
    FieldSchema("image_vector", dtype=DataType.BINARY_VECTOR, dim=64)
])

# 集合不存在，則會自動建立集合；已存在，不會重複建立
collection = Collection(settings.MILVUS_CONFIG.collection.name, schema)

使用的向量型別是 dtype=DataType.BINARY_VECTOR,

為什麼不選 float 是因為我不知道怎麼把 ahash 轉成 float

關於向量索引的問題，因為我選用的 『向量型別』 是 BINARY_VECTOR。所以，索引型別只有 BIN_FLAT 和 BIN_IVF_FLAT 可選了
具體可看 https://milvus.io/docs/v2.2.x...
圖片.png

插入 ahash 到 milvus

class TestVector(unittest.TestCase):
    def test_insert_vector(self):
        """
        插入 ahash 到 milvus
        python -m unittest testing.test_milvus.TestVector.test_insert_vector
        """

        oss_file_path = 'image_hash/testing/WechatIMG193.jpeg'

        file = BytesIO(bucket.get_object(oss_file_path).read())
        vector = create_vector(file)
        m_pk = insert_vector(vector, meta_id=2, company_id=1)
        logger.debug(f'milvus pk: {m_pk}')

查詢 ahash from milvus

def test_search(self):
    """
    批次呼叫後端介面入庫
    python -m unittest testing.test_milvus.TestVector.test_search
    """
    oss_file_path = 'image_hash/testing/WechatIMG193.jpeg'

    file = BytesIO(open(BASE_DIR/'testing'/'resource'/'WechatIMG193.jpeg','rb').read())
    vector = create_vector(file)

    logger.debug(vector)

    rows: list[dict[str, Any]] = collection.search(
        data=[vector],
        param={"metric_type": 'HAMMING', "params": {"nprobe": 32}},
        anns_field='image_vector',
        output_fields=['id', 'meta_id', 'company_id'],
        limit=10,
    )
    logger.debug(rows)
    logger.debug(type(rows))

注意 metric_type ，因為我選用的 『向量型別』 是 BINARY_VECTOR，所以，metric_type 要選擇支援 BINARY_VECTOR 的才行

圖片.png

相關文章