【Node.js】画像がえっちかAIで判定してみた

はじめに

ChatGPTをはじめとするAIの有用性が知れ渡り、検索エンジンや問い合わせフォーム、果ては手元にあるスマホ上など、ありとあらゆる場所への採用が検討されているこの大AI時代。火付け役こそChatGPTのような大規模言語モデルやStable Diffusionのような画像生成AIでしたが、世の中にはそれ以外にも様々なAIがあって、それぞれ複雑なタスクをこなすことができます。コンテンツモデレーションもその一つ。X（旧Twitter）などのSNSに投稿された画像は、本当に投稿・共有して大丈夫なものなのか、AIによって判別されていると言われています。

……で、実はMisskeyという日本発のオープンソースなSNS実装にも、AIを使って画像がセンシティブか（はたまたポルノか）を判定する機構が備わっています。MisskeyはNode.jsで動くのですが、この部分には、nsfwjsが使われていて、これは@tensorflow/tfjs-nodeを使っています。MisskeyはTensorFlowでAIを動かしているというわけです。すごいですね。

私は今、MisskeyをハードフォークしたSNS（Wisteria）を細々と開発していて、ちょうどこのnsfwjs周りをリファクタリングしようとしているところだったのですが……少し問題が発生してしまいました。なぜか最新のnsfwjsが私の環境では読み込めないようなのです。原因究明のためにnsfwjsのコードを読んだりしたのですが、CJS／ESMやビルドの仕方などなかなか複雑で、結局何が何やらといった感じでした。またそもそもnsfwjsは、最近はdependabotしかコミットしていない、実質メンテナンスされていないような状況だったため、いろいろ悲しくなりました。

というわけで、実際に使うかはともかくとして、nsfwjsに代わるいい感じのライブラリを作ってみることにしました。AIについてはまったくの素人で、この記事を書いている現在もまったく理解できていませんが、とりあえず動く形にはなったので備忘録としてここにいろいろと情報を残しておきます。

https://github.com/okayurisotto/echikana

要件定義と技術選定

まずは要件定義です。

• Node.jsで動くこと
• ESMを採用したモダンな開発環境でエラーなく読み込め、使えること
• TypeScriptの型定義があること（ライブラリ自体もTypeScriptで開発したい）
• 偽陽性や偽陰性が少ないこと（AIの性能が高いこと）

nsfwjsはWebブラウザなどのNode.js以外のランタイムでも動かすことができるのですが、AIのモデルファイルをクライアントに配信するのは流石にバイオレンスで需要がないので、私のライブラリでは非対応ということにしました。また、Cloudflare Workersなどのエッジランタイムも、使用可能なCPU時間の制限が厳しそうなことからひとまず考えないものとしました。BunやDenoなどにも対応しません。というかこのあたりはWebGPU APIの普及によって状況が比較的すぐに変わりそうなので……。

さて、ざっと調べてみただけでも、Node.jsで使えるAIのランタイムにはいくつか種類がありそうでした。nsfwjsでも採用されていた@tensorflow/tfjs-nodeはもちろん第一候補でした（過去形）。とりあえず@tensorflow/tfjs-nodeを採用する方針で考えてみていました（過去形）。

あとやはり大事なのはどういったAIモデルを採用するか、でしょう。nsfwjsではGitHubで配布されているGantMan/nsfw_modelが使われているようです。しかし今回、より偽陽性・偽陰性の少ない高性能なAIを使いたいということで、とりあえずこれとは別の選択肢がないか探してみました。すると、Hugging FaceにFalconsai/nsfw_image_detectionというものを見つけました。先月だけでも150万回近くダウンロードされている、なかなかすごそうなものです。ひとまずこれを動かせないか模索することにしました。

https://huggingface.co/Falconsai/nsfw_image_detection

動かない。なんもわからん。

AIモデルとひとくちに言っても、ランタイムごとにいろいろと事情があるのか、いくつかの形式があるようでした。@tensorflor/tfjs-nodeなどでは、tensorflowjs_converterというツールを使って、従来の様々な形式のモデルをtfjs向けモデルに変換できます。

https://github.com/tensorflow/tfjs/tree/master/tfjs-converter

……変換できるはずでした。インストールできませんでした。Dockerfileを書いて作ったまっさらなPython環境、もちろんバージョンはドキュメントの指定通りの3.6.8、pipを使ってtensorflowjs[wizard]をinstall……できませんでした。Issueを探してもよくわからず、そもそも「3.6.8よりずっと新しいバージョンのPython環境でエラーになった人に、それよりは古いものの3.6.8よりだいぶ新しいバージョンのPython環境の使用を勧める」やり取りが行われています。なんもわからん。……docker image build .コマンドの実行がホストOSのストレージの枯渇によって失敗したとき、私はTensorFlowを諦めました。

ONNX RuntimeとOptimum

AIランタイム第2候補はONNX Runtimeでした。

https://www.npmjs.com/package/onnxruntime-node

なんとあのMicrosoftが開発しています。TensorFlowがGoogle製なことを考えると、ライバル関係にあるのでしょうか？　npm trendsを見てみたところ、onnxruntime-nodeのダウンロード数は半年ほど前に@tensorflow/tfjs-nodeのダウンロード数を抜いたようです。

https://npmtrends.com/@tensorflow/tfjs-node-vs-onnxruntime-node

ONNX RuntimeはONNXというクロスポラットフォームな形式のAIモデルファイルを読み込み、実行します。既存の形式のモデルをこの形式に変換してくれるのはOptimumというツールです。Hugging Faceが開発しているようです。

https://github.com/huggingface/optimum

Optimumのメインの機能はOptimizeなようで、いろいろなオプションを受け取るようになっていますが、ただ変換したいだけならとても簡単です。変換先形式とHugging Faceのrepo/userで変換元モデルを指定するだけです。

$ optimum-cli export onnx --model Falconsai/nsfw_image_detection /path/to/model-dir

pipでのインストールも問題なくできました。

https://github.com/okayurisotto/echikana/blob/main/Dockerfile

https://github.com/okayurisotto/echikana/blob/main/compose.yaml

なぜかonnxruntime-nodeに型定義ファイルがない

モデル問題がなんとかなったので、あとは頑張ってプログラミングをすればいいだけです。……いいだけなはずでした。

https://github.com/microsoft/onnxruntime/issues/17979

はい、なぜかonnxruntime-nodeはTypeScriptの型定義ファイルを生成してくれていません。このライブラリの開発言語はTypeScriptですが、（そしてそもそもTypeScriptはMicrosoftの開発するものですが、）なぜかTypeScriptで使いづらいことになってしまっています。

上記Issueには、「ないなら作ればいいじゃん！」というような回避策が投稿されています。たしかにそうですね、ないなら作ればいいです。というわけで作りました。

https://github.com/okayurisotto/echikana/blob/main/patches/onnxruntime-node%401.18.0.patch

pnpmのpatch機能を使って、パッケージに型定義ファイルを挿入するようにしました。patch機能の存在は一応知ってはいましたが、まさか使う日が来るとは思ってもみませんでした……。

onnxruntimeの使い方

というわけで実装です。

onnxruntime（ort）ではInferenceSessionのメソッドを呼び出す形で推論モデルのロードや実行、アンロードを行います。InferenceSessionのcreateスタティックメソッドはArrayBufferLikeとしてモデル本体、もしくは文字列としてモデルへのパスを受け取って、InferenceSessionインスタンスを返すようになっています。そして、そのインスタンスのrunメソッドに、あらかじめ作っておいたTensorインスタンスを渡すことで推論が非同期的に行われます。返り値に含まれるTensorインスタンスを見れば、どういった推論が行われたのか確認できます。

/* とても大雑把な流れ */

import ort from "onnxruntime-node";

// モデルの読み込み
const model = fs.readFileSync("path/to/model.onnx");
const inferenceSession = await ort.InferenceSession.create(model);

// 入力するTensorの用意
const input = new Float32Array([]); // いい感じに用意しておく
const inputTensor = new ort.Tensor(input);

// 推論
const result = await inferenceSession.run({ pixel_values: inputTensor });

// 出力を確認
const outputTensor = result["logits"];
const output = outputTensor.data;
console.log(output); // 今回は要素数2の`Float32Array`が出力される

とてもシンプルですね。

Tensorを受け取ってTensorを返すというこれを感覚的に理解するときには、以前YouTubeで見た動画が役立ってくれました。有名ですのですでに知っている人も多いかもしれませんが、面白い動画ですのでぜひ。

https://www.youtube.com/watch?v=tc8RTtwvd5U

実装：前処理

ただし気をつけないといけないことがあります。それは入力するFloat32Arrayの準備です。入力する画像データはあらかじめリサイズしたり標準化したりといった、前処理をしておく必要があります。どういった前処理をすればいいかは、モデルのpreprocessor_config.jsonというファイルで宣言されています。

{
  "do_normalize": true,
  "do_rescale": true,
  "do_resize": true,
  "image_mean": [
    0.5,
    0.5,
    0.5
  ],
  "image_processor_type": "ViTImageProcessor",
  "image_std": [
    0.5,
    0.5,
    0.5
  ],
  "resample": 2,
  "rescale_factor": 0.00392156862745098,
  "size": {
    "height": 224,
    "width": 224
  }
}

これによると、

• 画像を縦横224ピクセルにリサイズする
• それぞれのピクセルの色を構成するRGBそれぞれを、0x00 ~ 0xffから-1.0 ~ +1.0にする

必要がありそうです。

また、Tensor化する前にtranspose（行列の転置）をしておく必要もあるらしいです。赤色だけ、緑色だけ、青色だけ抽出した画像データを用意して、それらを繋げる形になるようです。

https://onnxruntime.ai/docs/tutorials/web/classify-images-nextjs-github-template.html

このあたりの画像処理では、今回はsharpを使うことにしました。sharpでは、JPEGやPNGなどにエンコードされた画像データをArrayBufferLikeなどとして渡して、リサイズや生のピクセルデータの取得ができます。

import fs from "node:fs";
import ort from "onnxruntime-node";
import sharp from "sharp";

const image = fs.readFileSync("path/to/image.jpg");
const pixels = await sharp(image)
  .resize(224, 224, {
    kernel: "nearest",
    fit: "fill", // アスペクト比は無視していいはず
  })
  .removeAlpha() // アルファチャネルは先に除いておく
  .raw() // 生のピクセルデータを得る（RGBRGBRGB...と並んでいる）
  .toBuffer();

const redComponents: number[] = []; // 赤色のみ
const greenComponents: number[] = []; // 緑色のみ
const blueComponents: number[] = []; // 青色のみ

for (let i = 0; i < pixels.byteLength; i += 3) {
  const red = pixels[i + 0];
  if (red === undefined) throw new Error();
  redComponents.push(red);

  const green = pixels[i + 1];
  if (green === undefined) throw new Error();
  greenComponents.push(green);

  const blue = pixels[i + 2];
  if (blue === undefined) throw new Error();
  blueComponents.push(blue);
}

// 転置
const transposed = [
  ...redComponents,
  ...greenComponents,
  ...blueComponents,
];

// -1.0 ~ +1.0 の範囲に標準化
const normalized = transposed.map((v) => (v / 0xff - 0.5) / 0.5);

const inputTensor = new ort.Tensor(
  new Float32Array(normalized),
  [1, 3, 224, 224], // 1 * 3 * 224 * 224 === normalized.length
);

実装：後処理

今回のモデルは、「NSFWではない確信度」と「NSFWである確信度」の2種類の情報を、要素数2のFloat32Arrayとして出力します。単純にどちらであるかを知りたいだけなら、どちらの要素が大きいか比べればいいだけですね。今回は具体的な数値として知りたかったので、softmax関数を実装しました。

const softmax = (values: readonly number[]): number[] => {
  const denominator = values
    .map((value) => Math.exp(value))
    .reduce((prev, current) => prev + current, 0);

  return values.map((value) => Math.exp(value) / denominator);
};

私自身この式が何をしているのかよくわかっていませんがこの関数に結果を渡してあげることで、0から1までの数値として、確信度を得ることができます。

完成

そんな感じで完成しました。渡された画像がNSFWかどうか判定するライブラリ、名付けてEchikana。

https://github.com/okayurisotto/echikana

試しに私がSNSでアイコンとして使っている画像（AI生成）を渡してみると……

0.12188608139140207

どうやら私のアイコンはSafe for Workなようです。

ではインターネットで適当に見繕ったえっちなイラストを渡してみると……

0.9987144074498981

面白いくらいに大きな数値が出てきました。

あとはえっちイラストではなく実写画像を渡してみると……

普通の実写画像: 0.02400125250067131
実写ポルノ画像: 0.9999112206326956

実写・イラスト問わずきちんと判定できていそうですね。

GantMan/nsfw_modelとの比較はまだできていません。気が向いたときにやろうかと思います。

おわりに

GitHubリポジトリには変換済みモデルは置いていませんが、Dockerfileやcompose.yamlファイルとしてダウンロードしてきて変換する処理がまとめてあります。だいぶ試しやすく整備できたと思いますので、興味があればお気軽にどうぞ。気が向いたらパッケージ本体をnpmで配布してみようと思います（が、パッケージ自作はやったことがないので……）。

以上です。最後まで読んでいただきありがとうございました。質問や指摘等ありましたら遠慮せずにどうぞ。