FaissでKNN実装&Streamlitで超お手軽(手抜き)文字認識

# 概要
先日の[ブログ記事](https://qiita.com/anyai_corp/items/1d66feea6102c28dd077)に関連し、近似最近傍探索ライブラリ [Faiss](https://github.com/facebookresearch/faiss) による類似データ検索の実装例を紹介します。  
さらに、[Streamlit](https://streamlit.io/)を使って超お手軽(手抜き)文字認識アプリで遊んでみます。単純にMNIST画像のピクセル値を比較するだけですが、そこそこ判定できそうです（改めて深層学習のありがたみも感じました）。  
興味ある方、トライしてみてください！

# はじめに
類似データ検索は、商品検索や顔識別など広く活用されている技術です。  
検索対象から特徴量を抽出（ベクトル化）し、ベクトル同士の類似度をコサイン類似度やユークリッド距離などの指標に基づいて算出することにより、類似度の高いデータを検索できます。  
ベクトル化は深層学習モデルが得意としており、画像や文章など複雑なデータにおいても高度な特徴抽出が可能になってきました。

なお、類似データ検索ではベクトル化だけでなく**類似度計算の手法**も重要になります。  
特徴量の次元数をD、検索先のサンプル数をN、距離指標をユークリッド距離とした場合、最近傍探索では計算量はDとNに依存します。サンプル数が多い場合は処理時間を要するため、近年は検索精度を低下させる代わりに計算量を削減できる**近似**最近傍探索の利用が一般的です。  
代表的な近似最近傍探索のPythonライブラリには、Faissやnmslibなどがあり、今回はFaissを使用します。


# 検証条件
- 環境、ライブラリ
  - python : 3.10.11
  - faiss-cpu==1.7.4
  - keras==2.13.1
  - matplotlib==3.8.0
  - numpy==1.24.3
  - Pillow==9.5.0
  - streamlit==1.27.0
  - streamlit-drawable-canvas==0.9.3
  - tensorflow==2.13.0


# データセットの用意
MNISTデータで動作確認します。
```python
from typing import Union

import numpy as np
import matplotlib.pyplot as plt
from keras.datasets import mnist
import faiss

# MNISTデータ読込
(X_train, y_train), (X_test, y_test) = mnist.load_data()
X_train.shape, X_test.shape
>>> ((60000, 28, 28), (10000, 28, 28))

# 可視化
fig, axs = plt.subplots(4,4, figsize=(8, 8))
for i, (image, label) in enumerate(zip(X_train, y_train)):
    axs[i//4][i%4].imshow(image)
    axs[i//4][i%4].set_title(label)
    if i == 15:
        break
plt.tight_layout()
plt.show()
```
![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3476911/8b138499-9155-7f5f-5f92-2ef34f68cf89.png)


# K近傍探索の実装

Faissは検索先のデータを `index` に登録し、`search`メソッドにより近傍探索を実施します。  
前処理など含め scikit-learnライクに実行出来るよう、`fit`・`predict`メソッドでラップし、  
引数 `k`で探索数、`metric`で距離指標を`euclid(ユークリッド距離)`か`cosine(コサイン距離)`に設定します。

```python
class FaissKNeighbors:
    def __init__(self, k: int = 20, metric: Union["euclid", "cosine"] = "euclid"):
        self.index = None
        self.d = None
        self.k = k
        self.metric = metric

    def fit(self, X: np.ndarray):
        X = X.copy(order="C")
        self.d = X.shape[1]
        X = X.astype(np.float32)
        if self.metric == "cosine":
            self.index = faiss.IndexFlatIP(self.d)  # cosine
            faiss.normalize_L2(X)
        elif self.metric == "euclid":
            self.index = faiss.IndexFlatL2(self.d)  # euclid
        self.index.add(X)

    def predict(self, X: np.ndarray):
        X = X.copy(order="C")
        X = np.reshape(X, (-1, self.d))
        X = X.astype(np.float32)
        if self.metric == "cosine":
            faiss.normalize_L2(X)
        distances, indices = self.index.search(X, k=self.k)
        if self.metric == "euclid":
            distances = np.sqrt(distances)
        if X.shape[0] == 1:
            return distances[0], indices[0]
        else:
            return distances, indices
```

# 動作確認
```python
knn = FaissKNeighbors(k=5, metric='euclid')
# X_trainデータを検索先として登録
knn.fit(X_train.reshape(-1, 28*28))
# X_testデータ100件分 検索
pred_dists, pred_indexs = knn.predict(X_test[:100].reshape(-1, 28*28))

# 検索結果の可視化
for i in range(5):
    fig, axs = plt.subplots(1,6)
    axs[0].imshow(X_test[i])
    axs[0].set_title(f'Target : {y_test[i]}')
    for col, p_i in enumerate(pred_indexs[i]):
        axs[col+1].imshow(X_train[p_i])
        axs[col+1].set_title(f'pred{col+1} : {y_train[p_i]}')
    plt.tight_layout()
    plt.show()
```
![image-1.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3476911/a5d60c90-252c-1263-c28f-8964abbaac77.png)


検索対象(Target)と類似した画像が検索できています(pred1~5)。


# おまけ
MNISTは手書き文字データセットなので、実際に自分の手書き文字も検索・判定出来るか気になってきました。  
そこで、ChatGPTに相談しつつ手書き文字認識アプリをStreamlitで実装してみました。文字入力はStreamlit拡張機能の[streamlit-drawable-canvas](https://github.com/andfanilo/streamlit-drawable-canvas)を利用しています。数字の判定は、類似度上位5件の多数決で決めてます。

前述のFaissKNeighborクラスをimport出来るよう、faiss_kneighbors.pyとしてスクリプト化
```faiss_kneighbors.py
from typing import Union

import numpy as np
import faiss


class FaissKNeighbors:
    def __init__(self, k: int = 20, metric: Union["euclid", "cosine"] = "euclid"):
        # 上記「K近傍探索の実装」と同様のため割愛
        # ...
```
アプリ情報をdemo.pyに実装
```demo.py
import faiss
import numpy as np
from PIL import Image
from keras.datasets import mnist
import streamlit as st
from streamlit_drawable_canvas import st_canvas

from faiss_kneighbors import FaissKNeighbors

K_NUM = 5

_, (X_test, y_test) = mnist.load_data()
knn = FaissKNeighbors(k=K_NUM)
knn.fit(X_test.reshape(-1, 28*28))

# Streamlitアプリケーションの開始
st.title("手書き文字認識アプリ")

# マウスで文字を描画するキャンバスを作成
st.write("マウスで文字を描いてください。")
canvas_result = st_canvas(
    stroke_width=20,
    update_streamlit=False,
    height=200,
    width=200,
    drawing_mode='freedraw',
)

# 描画した文字を整形し、KNNを使用して分類
if canvas_result.image_data is not None:
    drawn_image = canvas_result.image_data
    drawn_image_gray = drawn_image[:, :, 3]
    if np.sum(drawn_image_gray) > 0:
        drawn_image_gray = Image.fromarray(drawn_image_gray)
        resized_image = drawn_image_gray.resize((28, 28))
        resized_image = np.array(resized_image)

        # KNNで分類
        dists, indexs = knn.predict(resized_image)
        votes = y_test[indexs]
        predictions = np.flatnonzero(np.bincount(votes) == np.bincount(votes).max())

        st.write(f"判定結果: {predictions}")
        st.write(f"類似度上位")
        cols = st.columns(K_NUM)
        for i, idx in enumerate(indexs):
            pred_image = Image.fromarray(X_test[idx])
            pred_image = pred_image.resize((100,100))
            pred_image = np.array(pred_image)
            cols[i].image(pred_image, clamp=True, caption=f'label = {y_test[idx]}')
```
アプリ起動
```
streamlit run demo.py
```

こんな感じで動きます。  
単純な仕組みなため、4、5、7、9は他の数字と間違えてしまってますね...。  
数字を書く位置にとても敏感なため、実用化は厳しいですね。CNNの位置不変性が活躍してくれそうです。
![output.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/3476911/f38e55db-8821-0909-7d56-285df0ac8acd.gif)


最後まで読んでいただき、ありがとうございました！  
今後も機械学習の活用を始め、開発環境やシミュレーションなど幅広く技術情報発信をしていく予定です！  

最後になりますが、本記事の内容に誤りなどあれば、コメントにてご教授お願いいたします。