이미지 Embedding 하여 VectorDB에 저장하기

이번에는 Image 를 Embedding 해서 VectorDB에 저장하고, 새로운 이미지에 대해 가장 유사한 이미지를 찾아보는 실험을 해보도록 하겠습니다.

Image Embedding 방법은 imagenet으로 pre-trained 된 EfficientNetV2B1 모델을 FC(Top Layer)를 제외하고 사용하였고, Global Average Pooling 2D를 추가해주었습니다.
최종 Layer가 1280차원이며, 이 값을 VectorDB에 저장을 하겠습니다.

VectorDB는 pgvector를 사용합니다.

시나리오는 아래와 같습니다.

1. imagenet-1k 의 validation 데이터 5만장을 이미지 Embedding을 합니다.
2. Embedding 된 데이터를 pgvector에 저장합니다.
3. Embedding에 사용되지 않은 이미지를 Vector검색을 하고, 가장 유사한 이미지 3장을 출력해봅니다.

---

이미지 Embedding 및 저장 수행

먼저, DB에 Table을 아래와 같이 생성하였습니다.

CREATE TABLE imagenet(id bigserial PRIMARY KEY, filename varchar(30), embedding vector(1280))

그리고 Embedding에 사용하는 모델은 Tensorflow를 이용하였고, VectorDB에 저장까지 하는 코드는 아래와 같습니다.

import os
import numpy as np
import tensorflow as tf
import psycopg2

from tensorflow.keras.applications.efficientnet_v2 import EfficientNetV2B1, preprocess_input

from glob import glob

BATCH_SIZE = 32
INPUT_SHAPE = (224, 224)

image_paths = glob('D:/data/imagenet1k/Data/CLS-LOC/val/*.jpeg')
print(len(image_paths))
# 50000

@tf.function
def load_image(path):
    image = tf.io.read_file(path)
    image = tf.image.decode_jpeg(image, channels=3)
    image = tf.image.resize(image, INPUT_SHAPE)  
    image = image / 255.0
    return image, path
    

# Build DataLoader
AUTOTUNE = tf.data.experimental.AUTOTUNE

ds = tf.data.Dataset.from_tensor_slices(image_paths)
ds = ds.map(load_image, num_parallel_calls=AUTOTUNE)
ds = ds.batch(BATCH_SIZE).prefetch(buffer_size=AUTOTUNE)

# Build Model
inputs = tf.keras.layers.Input(shape=INPUT_SHAPE + (3,))
feature_extractor = EfficientNetV2B1(include_top=False, weights='imagenet')(inputs)
outputs = tf.keras.layers.GlobalAveragePooling2D()(feature_extractor)

model = tf.keras.Model(inputs=inputs, outputs=outputs)
model.summary()

# Embedding & Insert Data to pgvector
conn = psycopg2.connect('postgresql://namkon:${SECRET}@IP_ADDRESS:5432/test')
cursor = conn.cursor()
sql = "insert into imagenet (filename, embedding) values ('{filename}', '{embedding}')"

for images, paths in ds:
    embeddings = model(images)
    for embedding, path in zip(embeddings.numpy(), paths.numpy()):
        filename = os.path.basename(path.decode())
        cursor.execute(sql.format(filename=filename, embedding=list(embedding)))

conn.commit()

위와 같이 실행한 결과, 아래와 같이 pgvector에 저장이 잘 된것을 확인했습니다.

---

결과 테스트

이제 저장된 Embedding 값 기반으로, 새로운 이미지를 가지고 유사한 이미지를 찾아보고 얼마나 비슷한지 확인을 해보도록 하겠습니다.
이미지 샘플은 임의로 선택을 했고, 코드와 결과는 아래와 같습니다.

test_path = 'D:/data/imagenet1k/Data/CLS-LOC/test/ILSVRC2012_test_00001580.jpeg'
test_image = tf.keras.utils.load_img(test_path, target_size=(224, 224))
test_image = tf.keras.utils.img_to_array(test_image)
test_image = test_image[tf.newaxis, ...] / 255.

embedding = model(test_image)

%%time
sql = f"select * from imagenet order by embedding <-> '{list(embedding.numpy()[0])}' limit 3"
cursor.execute(sql)
ret = cursor.fetchall()
# CPU times: total: 0 ns
# Wall time: 297 ms

base_path = 'D:/data/imagenet1k/Data/CLS-LOC/val'

plt.subplots(figsize=(20,12))

plt.subplot(1,4,1)
plt.imshow(test_image[0])
plt.axis('off')
plt.title('Test Image')

plt.subplot(1,4,2)
plt.imshow(Image.open(os.path.join(base_path, ret[0][1])))
plt.axis('off')
plt.title('1st close image')

plt.subplot(1,4,3)
plt.imshow(Image.open(os.path.join(base_path, ret[1][1])))
plt.axis('off')
plt.title('2nd closest image')

plt.subplot(1,4,4)
plt.imshow(Image.open(os.path.join(base_path, ret[2][1])))
plt.axis('off')
plt.title('3rd closest image')

plt.show()

---

Conclusion

그렇습니다.. 5만개 데이터를 벡터서치하는데 300ms 밖에 안걸렸다는 건 엄청 빠르다고 생각이 되지만, 결과가 개떡같이 나왔습니다 🥲

imagenet 1k 로 pre-trained 된 efficientNetV2B1 모델의 성능 문제일까요?

코드상에 문제가 없는 지 한번 더 점검을 해보고.. 이상이 없다면,
다음번엔 imagenet 21k로 Pre-trained & imagenet 1k로 fine-tuning된 Swin-TransformerV2 모델을 Embedding 모델로 테스트와 CLIP 모델을 통한 Embedding 을 해보도록 하겠습니다.