👙

けしからん画像分類器を作ってみる (10) 公開

2022/11/16に公開

目次

「けしからん画像分類器」を公開しました

前回の記事からかなりの時間が空いてしまいましたが、Twitterでモデル公開のリクエストを頂いたので久しぶりに再学習し、分類モデルを公開しました。

https://twitter.com/discord_nana/status/1587419833653932032

ONNXモデル、サンプルコード、ライセンス情報を含むZIPファイルは、以下のURL(Googleドライブ)からダウンロードできます。
なお、ZIPファイルのままであれば再配布して頂いて構いません。

ONNXモデル、サンプルコードはMITライセンスとしています。機械学習モデルの公開に際して、より適切なライセンスがあれば情報を頂けると幸いです。

実行例

ZIPファイルにはONNXモデルの他、推論のサンプルコードpredict.pyが含まれています。実行例を以下に示します。

pip3 install -r requirements.txt
python3 predict.py foo.jpg

predict.pyの内容は以下の通りです。

predict.py
#!/usr/bin/env python3

import os
import sys

import numpy as np
import onnxruntime
import PIL.Image

ONNX_MODEL_FILE = os.path.join(os.path.dirname(__file__), "kleamp1e-classifier-pornography-20221113_164725.onnx")
TARGET_SIZE = 384


def main():
    image_path = sys.argv[1]

    session = onnxruntime.InferenceSession(
        ONNX_MODEL_FILE, providers=["CPUExecutionProvider"]
    )
    session.get_modelmeta()
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    image = PIL.Image.open(image_path)
    image = image.resize(
        (TARGET_SIZE, TARGET_SIZE), resample=PIL.Image.Resampling.BICUBIC
    )
    image = image.convert("RGB")
    image = np.array(image, dtype=np.float32) / 255
    image = np.expand_dims(image, 0)

    predictions = session.run([output_name], {input_name: image})
    print(predictions[0][0][0])


if __name__ == "__main__":
    main()

分類モデルについて

分類モデルの概要は以下の通りです。

バックボーンネットワーク

以前の記事では、バックボーンネットワークとして「EfficientNet B0」を使用していましたが、今回のモデルは「EfficientNet V2 S」を使用しています。

入力

入力は画像で、サイズは(1, 384, 384, 3)、並びはNCHWです。リサイズ以外の前処理は必要ありません。

出力

画像が「けしからん」度合いを0から1の範囲で出力します。1が「けしからん画像」です。

データ数

テストデータ、バリデーションデータ、学習データのデータ数は以下の通りです。なお、学習時における不均衡データに対する対処は(このモデルでは)行っていません。

種別 ラベル:0 ラベル:1 合計
テストデータ 1002 1002 2004
バリデーションデータ 1036 1036 2072
学習データ 8343 13326 21669
合計 10381 15364 25745

混同行列

テストデータの混同行列(Confusion Matrix)は以下の通りです。

推論結果:0 推論結果:1
ラベル:0 882 120
ラベル:1 60 942

メトリクス

テストデータの各種メトリクスは以下の通りです。

メトリクス 数値
正解率(Accuracy) 0.9101
適合率(Precision) 0.8870
再現率(Recall) 0.9401
F1スコア 0.9127
AUCスコア 0.9745

ROC曲線

テストデータのROC曲線は下図の通りです。

定性評価

定性評価として、画像と推論結果をざっと眺めてみました。

今回のデータセットには写真(3次元)とイラストなどの絵(2次元)の両方を含んでいますが、イラストについては「けしからん」と誤分類する傾向にあるようです。
データセットに含まれるイラストには「けしからん」ものが多く、そうでないデータが少ないことが原因である可能性があります。

写真については概ね正しく分類できているようですが、「けしからん」と判定される要素が画像の面積的に小さい場合、正しく判定できていないことが多いように感じました。
これについては、モデルの入力が384x384ピクセルと比較的低い解像度のため、前処理のリサイズの時点で情報が欠如していることが原因である可能性があります。

おわりに

随分前に作った分類モデルではありますが、今回の公開にあたり再学習と評価を行いました。
評価結果をまとめることで「分類モデルを公開(出荷)する際には、どのような情報が必要か」について学ぶことができました。

この記事や分類モデルが、何かのお役に立てば幸いです。

Discussion