NVIDIAのAnomalyGenで、良品写真から不良品画像を合成してみた
はじめに
以前の記事では、NVIDIAのCosmos3を用いて欠陥品の合成画像を作ろうというトライをして、思うようにいきませんでした。
具体的には、生成される画像がCGっぽいものになってしまったり、既存の実写良品画像を基に欠陥画像を生成するのが難しいといった課題がありました。
他の生成方法を探しているうちに見つけたものが、今回検証するNVIDIAのAnomalyGenです。
良品画像＋「ここに欠陥を入れたい」というマスク＋欠陥の種類名を渡すと、マスクの中にリアルな欠陥を描き込みしてくれる、少数ショット前提の合成データ生成パイプラインです。
今回はこれをAWSの環境で使い、磁性タイルの不良データ生成までやってみました。
先に結論を書くと、
- 学習（post-training）約1時間＋生成3.2秒/枚で、良品写真に「欠け」「気孔」「ひび」「研磨ムラ」を描き込めました。マスクの外は元の良品写真のまま保たれます。
- 学習といっても動かすのは全体2.6Bパラメータのうちわずか340万（0.13%）。ベースの生成モデルは凍結されていて、だから1GPUで1時間で終わります。
- ただし公式の手順どおりのやり方一発では動きませんでした。詰まりどころが4つあったので、そこを中心に紹介します。
AnomalyGenとは
- リポジトリ: NVIDIA/paidf-anomalygen
- 公式の説明によると「少数のサンプル画像を基に、異常画像の合成データ生成（SDG）のための拡散モデルベースのパイプライン」。
- 仕組みとしては、凍結したCosmos-Predict2-2B-Text2Image（画像生成モデル）の上に、マスクエンコーダ（NV-DINOv2、これも凍結）と小さなアダプタ、そして欠陥の種類ごとの埋め込み（anomaly embedding）を乗せた構成です。実行ログで確認すると、学習で更新されるのはアダプタと埋め込みだけでした。
- 上記のリポジトリでは、リファレンスのユースケースが3つ同梱されています: UC1=プリント基板、UC2=金属質の表面（公式カテゴリ名はMetal Surface。題材は磁性タイル）、UC3=スマホ画面ガラス。今回はUC2でやります。
UC2のサンプルデータは、公開されている磁性タイルの欠陥データセット（abin24/Magnetic-tile-defect-datasets）から準備されます（python3 -m scripts.utilities.prepare_dataset_uc2 <output_dir>で自動ダウンロードされる）。このデータセットには5種の欠陥画像（ピクセル単位の正解マスク付き）に加えて、欠陥なしの良品画像（MT_Free）も含まれていて、準備スクリプトはここから良品と欠陥サンプルの両方を取り出します。欠陥クラスはMT_Blowhole（気孔）、MT_Break（欠け）、MT_Crack（ひび）、MT_Fray（ほつれ）、MT_Uneven（研磨ムラ）の5種類を使います。
学習に使う実欠陥画像は1タイプあたりたった5枚（＋マスク5枚）、良品20枚という構成です。「不良品の実画像が数枚しかない」という現場の状況でも利用可能です。
環境
| 項目 | 内容 |
|---|---|
| インスタンス | g6e.2xlarge（NVIDIA L40S 48GB×1、8vCPU、RAM 64GB） |
| AMI | Deep Learning Base OSS Nvidia Driver GPU AMI（Ubuntu 24.04） |
| 実行環境 | 公式のDockerイメージ（developモード） |
最初は一段安いg6e.xlarge（RAM 32GB）で始めたのですが、これが最初の詰まりどころになりました。
全体の流れ
基本はリポジトリのREADMEどおりです。
ざっくりの流れはこうです。
- Dockerイメージをビルドして起動する（
build_image.sh）。 - コンテナ内でベースモデルの重みをダウンロードする（
scripts.download_checkpoints）。 - サンプルデータセットを準備する（UC2は
prepare_dataset_uc2）。 - 学習用のconfig（YAML）を用意して学習（post-training）を回す。
- 生成指示のjsonl（testcase）を作って一括生成する。
この流れ自体は公式どおりなので、この記事では詰まった4つのポイントを中心に書いていきます。
ポイント1: DockerビルドでインスタンスごとOOM → RAM 64GBに変更
g6e.xlarge（RAM 32GB）でdevelopイメージをビルドしたところ、flash-attentionのソースコンパイルの途中でSSHが切断され、インスタンスのステータスチェックがimpairedになりました。ホストのメモリを食い尽くしたようです。
対処はインスタンスタイプの変更です。EBSはそのまま、stop→タイプ変更→startでg6e.2xlarge（RAM 64GB）にしたら、ビルドは通りました。Dockerのビルドキャッシュもボリュームに残っているので、途中からやり直せます。
aws ec2 stop-instances --instance-ids <instance-id>
aws ec2 modify-instance-attribute --instance-id <instance-id> \
--instance-type "{\"Value\":\"g6e.2xlarge\"}"
aws ec2 start-instances --instance-ids <instance-id>
GPUは同じL40S 1枚なので、増えたのはRAMとvCPUだけ。「GPUは足りてるのにホストRAMで死ぬ」パターンは、ソースからのビルドがあるリポジトリだと起きがちです。
なお、ベースモデルの重みダウンロードでも小さな想定外が2つありました。コンテナを起動し直すとHuggingFaceのログインが消えていること（hf auth loginのやり直し）と、hf側のxetという転送機構のエラー（Unable to parse string as hex hash value）です。後者はexport HF_HUB_DISABLE_XET=1で回避できました。
ポイント2: 学習が同梱のvalidation.jsonlで落ちる（PCB前提になっている）
金属用のconfigを用意して学習を回すと、いきなりこれで落ちました。
FileNotFoundError: [Errno 2] No such file or directory:
'datasets/MeiweiPCB/train_downsample/PCB/mask/defect/AAAAAA0823-...-Cur_mask.jpg'
学習configのdataloader_valが参照するag_inference/validation.jsonlが、リポジトリ同梱のPCBサンプル前提で、UC2のデータでは存在しないパスを読みに行きます。しかもtrainer.run_validation: Falseにしても回避できません（バリデーション用データセット自体は必ず初期化されるため）。
対処は、金属用の生成指示jsonlをcreate_testcaseで作って、configのinput_data_pathを差し替えることです。このjsonlはあとの生成ステップでも使うので、どのみち必要になります。
python -m scripts.anomaly_gen.create_testcase \
--OK_image_path datasets/UC2_data/metal_surface/clean_image \
--NG_mask_path datasets/UC2_data \
--name UC2_metal
ここにも小さな罠が2つありました。--NG_mask_pathは{テクスチャ}/mask/{欠陥種}/という階層の1つ上を渡す必要があること（スクリプトのソースを読んで確認しました）。もう1つ、探索対象のディレクトリ直下にdefect_spec.jsonlというファイルがあると、それをテクスチャ名として拾ってしまいNotADirectoryErrorになること。ファイルを一時的に退避して回避しました。
ポイント3: マスクは「厳密に0/255の2値」でないと弾かれる
次はこれです。
ValueError: Mask is not binary. Please check the mask format.
AnomalyGenのマスクチェックはソースを見るとnp.all(np.isin(mask, [0, 255]))、つまり0と255以外の画素が1つでもあるとアウトという厳密判定です。ところが磁性タイルデータセットの正解マスクは、エッジがアンチエイリアスされていて、画素値を調べると0と255のほかに13階調の中間グレーが混ざっていました。
対処は閾値127で2値化して上書きです。
import glob, numpy as np
from PIL import Image
for p in glob.glob('datasets/UC2_data/metal_surface/mask/**/*.png', recursive=True):
im = np.array(Image.open(p).convert('L'))
Image.fromarray(np.where(im > 127, 255, 0).astype('uint8')).save(p, format='PNG')
自社データでやる場合も、マスクは完全な2値で用意する必要がある、というのは覚えておいて損がないポイントだと思います。
ポイント4: 生成時に「画像とマスクのサイズ不一致」
学習は通ったのに、生成では今度はこれで落ちました。
AssertionError: Error: Image filename ... 's size with mask filename ...
磁性タイルのデータセットは画像ごとに解像度がバラバラで、create_testcaseは良品とマスクをランダムに組み合わせる（既定でマスク1枚あたり16件生成する）ので、サイズの合わないペアがほぼ確実に混ざります。リファレンスのPCBデータのディレクトリ名がtrain_downsample（＝サイズ統一済み）になっているのは、たぶんこのためです。
対処として、生成用に良品もマスクも512×512に統一したコピーを作りました（モデルの出力も512×512です）。マスクの縮小はNEAREST補間にして、ポイント3の2値を壊さないようにします。
結果: 学習前後のビフォーアフター
まず試運転として、リファレンスconfigのデフォルト値である200イテレーションだけ学習して、MT_Breakの5枚を生成してみました。結果は、肝心の欠陥がほぼ描き込まれませんでした。質感がうっすら変わる程度です。
元画像（良品画像）
200イテレーション画像
そこで14,000イテレーション（READMEに載っているUC1の学習量を参照。UC2での適正値は明記されていません）まで学習し
て同じ生成をすると、はっきり変わりました。
14,000イテレーション画像
ちなみに各欠陥クラスの様子は以下の通りでした。
- MT_Break: タイルの縁に欠けが生成される
- MT_Blowhole: 小さな気孔が出る
- MT_Crack: 細いひび。控えめだが指定位置に出る
- MT_Uneven: 大面積の研磨ムラが大胆に描かれる
マスクの外は元画像のままなので、「この良品写真の、ここに、この種類の不良」というラベル付きの不良データがそのまま手に入ります。出力にはTAO DAFT v3形式への変換も同梱されていて、下流の検査モデル学習にも渡しやすい形になっていました。
数字まわり
| 項目 | 実測値 |
|---|---|
| 学習速度 | 0.24秒/iter（batch=2、512×512）※初回のみCUDA Graph構築などで約9分 |
| 学習14000 iter | 約65分（うち初回イテレーションの約9分を含む） |
| 生成 | 400枚を1269秒＝約3.2秒/枚（35ステップ）※モデル初期化に別途約6分 |
| 学習対象パラメータ | 3,409,920（全体2.6Bの約0.13%） |
| VRAMピーク | 約25.6GB（学習〜生成を通した観測値） |
コスト
- g6e.2xlarge（東京）は、私が使用した際は$3.25168/時でした。学習＋全件生成の本番チェーンは90分ほどで、$5弱でした。環境構築の試行錯誤も含めたトータルでも$15前後です。
- 今回は夜間に無人で回したかったので、「学習→生成→結果をS3にアップロード→成功したら自動でstop」までを
&&で繋いだうえ、保険としてsudo shutdown -h +180を仕掛けておきました。GPUインスタンスの消し忘れ対策としておすすめです。
おわりに
次回は、このパイプラインをネット遮断環境（air-gapped）で動かすのを試します。工場のオフライン環境で完結できるかどうか、は製造業だと結構重要なポイントなので。
Built on NVIDIA Cosmos. ベースモデルとしてnvidia/Cosmos-Predict2-2B-Text2Image（NVIDIA Open Model License）を使用しています。
