PyTorchを使わず画像分類をする (imgbeddings)

使用するライブラリ

https://github.com/minimaxir/imgbeddings

READMEに記載がありますが以下の特徴があります。

• PyTorchやTensorFlowを必要としない
• 次元を削減するユーティリティが含まれる

今回、Macで手軽に画像検索を作りたかったため使ってみました。

画像の読み込みと埋め込み

画像の枚数 648枚で試して MacBook Pro (M2/32GB) で80秒ほどかかりました。
埋め込み後の次元は768次元です。

import os
from PIL import Image
from imgbeddings import imgbeddings

ibed = imgbeddings()

images = []
for file in os.listdir(IMAGE_DIR):
    image = Image.open(os.path.join(IMAGE_DIR, file))
    images.append(image)

embeddings = ibed.to_embeddings(images)
print(embeddings.shape)
# (<画像の数>, 768)

付属のユーティリティで次元削減して可視化

imgbeddingsに付属している次元削減を使用して可視化してみます。

from sklearn.cluster import KMeans
import matplotlib.pyplot as plt

ibed.pca_fit(embeddings, 2)
embeddings_pca = ibed.pca_transform(embeddings)

kmeans = KMeans(n_clusters=36, random_state=0)
clusters = kmeans.fit_predict(embeddings_pca)

sns.scatterplot(x=embeddings_pca[:, 0], y=embeddings_pca[:, 1], hue=clusters)
plt.show()

t-SNEで次元削減して可視化

比較のためにt-SNEを使用して可視化してみます。

from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

tsne = TSNE(n_components=2, random_state=0)
embeddings_2d = tsne.fit_transform(embeddings)

kmeans = KMeans(n_clusters=36, random_state=0)
clusters = kmeans.fit_predict(embeddings_2d)

sns.scatterplot(x=embeddings_2d[:, 0], y=embeddings_2d[:, 1], hue=clusters)
plt.show()

まとめ

普段はMacを使用しているためそのまま画像分類のモデルが動かせて便利でした。
(PyTorchやTensorFlowがなくても動くのは助かる)