latent diffusion segmentation
1.0.0
このリポジトリには、パノプティック セグメンテーションとマスク修復のためのシンプルな潜在拡散アプローチである LDMSeg の Pytorch 実装が含まれています。提供されたコードにはトレーニングと評価の両方が含まれています。
パノプティックセグメンテーションとマスク修復のためのシンプルな潜在拡散アプローチ
ウーター・ヴァン・ガンスベーケとバート・デ・ブラバンデール
この論文では、パノプティック セグメンテーションのタスクに取り組むための条件付き潜在拡散アプローチを紹介します。目的は、特殊なアーキテクチャ (例: 領域提案ネットワークやオブジェクト クエリ)、複雑な損失関数 (例: ハンガリアン マッチングまたはバウンディング ボックスに基づく)、および追加の後処理方法 (例: クラスタリング、NMS) の必要性を排除することです。 、またはオブジェクトの貼り付け)。結果として、私たちはタスクに依存しないフレームワークである Stable Diffusion に依存しています。提示されたアプローチは 2 つのステップで構成されます。(1) 浅い自動エンコーダを使用してパノプティック セグメンテーション マスクを潜在空間に投影します。 (2) RGB 画像を条件として、潜在空間で拡散モデルをトレーニングします。
主な貢献: 私たちの貢献は 3 つあります。
コードは、最新の Pytorch バージョン (2.0 など) で実行されます。さらに、Anaconda を使用して Python 環境を作成できます。
conda create -n LDMSeg python=3.11
conda activate LDMSeg
自動インストールに従うことをお勧めします ( tools/scripts/install_env.shを参照)。次のコマンドを実行して、プロジェクトを編集可能モードでインストールします。すべての依存関係が自動的にインストールされることに注意してください。これは常に機能するとは限らないため (CUDA や gcc の問題などにより)、手動インストール手順を参照してください。
python -m pip install -e .
pip install git+https://github.com/facebookresearch/detectron2.git
pip install git+https://github.com/cocodataset/panopticapi.git最も重要なパッケージは、次のように pip を使用してすばやくインストールできます。
pip install torch torchvision einops # Main framework
pip install diffusers transformers xformers accelerate timm # For using pretrained models
pip install scipy opencv-python # For augmentations or loss
pip install pyyaml easydict hydra-core # For using config files
pip install termcolor wandb # For printing and logging私の環境のコピーについては、 data/environment.ymlを参照してください。また、detectron2 と panopticapi からの依存関係の一部にも依存しています。彼らのドキュメントに従ってください。
現在、COCO データセットをサポートしています。画像とそれに対応するパノプティック セグメンテーション マスクをインストールするにはドキュメントに従ってください。また、COCO データセットに関するいくつかの例については、 ldmseg/data/ディレクトリを参照してください。補足として、採用される構造はかなり標準的なものである必要があります。
.
└── coco
├── annotations
├── panoptic_semseg_train2017
├── panoptic_semseg_val2017
├── panoptic_train2017 -> annotations/panoptic_train2017
├── panoptic_val2017 -> annotations/panoptic_val2017
├── test2017
├── train2017
└── val2017
最後になりましたが、 configs/env/root_paths.yml内のパスをデータセット ルートと目的の出力ディレクトリにそれぞれ変更します。
提示されたアプローチは 2 つの側面から構成されています。まず、低次元空間 (64x64 など) でセグメンテーション マップを表すように自動エンコーダーをトレーニングします。次に、事前トレーニング済みの潜在拡散モデル (LDM)、特に安定拡散から開始して、RGB 画像からパノプティック マスクを生成できるモデルをトレーニングします。モデルは、次のコマンドを実行してトレーニングできます。デフォルトでは、 tools/configs/base/base.yamlで定義されたベース構成ファイルを使用して COCO データセットでトレーニングします。 hydraパッケージに依存しているため、このファイルは自動的にロードされることに注意してください。
python - W ignore tools / main_ae . py
datasets = coco
base . train_kwargs . fp16 = True
base . optimizer_name = adamw
base . optimizer_kwargs . lr = 1e-4
base . optimizer_kwargs . weight_decay = 0.05引数の受け渡しの詳細については、 tools/scripts/train_ae.shを参照してください。たとえば、このモデルを 23 GB の単一 GPU で合計バッチ サイズ 16 で 50k 反復実行します。
python - W ignore tools / main_ldm . py
datasets = coco
base . train_kwargs . gradient_checkpointing = True
base . train_kwargs . fp16 = True
base . train_kwargs . weight_dtype = float16
base . optimizer_zero_redundancy = True
base . optimizer_name = adamw
base . optimizer_kwargs . lr = 1e-4
base . optimizer_kwargs . weight_decay = 0.05
base . scheduler_kwargs . weight = 'max_clamp_snr'
base . vae_model_kwargs . pretrained_path = '$AE_MODEL' $AE_MODEL前のステップで取得したモデルへのパスを示します。引数の受け渡しの詳細については、 tools/scripts/train_diffusion.shを参照してください。たとえば、このモデルを 16 GB の 8 GPU で合計バッチ サイズ 256 で 200k 反復実行しました。
いくつかのトレーニング済みモデルをリリースする予定です。 (クラスに依存しない) PQ メトリックは、COCO 検証セットで提供されます。
| モデル | #Params | データセット | イター | PQ | スクエア | RQ | ダウンロードリンク |
|---|---|---|---|---|---|---|---|
| AE | ～2M | ココ | 66k | - | - | - | ダウンロード（23MB） |
| LDM | ～800M | ココ | 200k | 51.7 | 82.0 | 63.0 | ダウンロード（3.3GB） |
注: 追加の微調整を実行しないため、強力でない AE (つまり、レイヤーのダウンサンプリングまたはアップサンプリングが少ない) は、多くの場合修復に役立ちます。
評価は次のようになります。
python - W ignore tools / main_ldm . py
datasets = coco
base . sampling_kwargs . num_inference_steps = 50
base . eval_only = True
base . load_path = $ PRETRAINED_MODEL_PATH 必要に応じてパラメータを追加できます。 --base.eval_kwargs.count_th 700や--base.eval_kwargs.mask_th 0.9などのしきい値を高くすると、数値がさらに増加する可能性があります。ただし、閾値を0.5とし、面積が512未満のセグメントを除いた標準値を用いて評価を行っています。
上記の事前トレーニング済みモデルを評価するには、 tools/scripts/eval.shを実行します。
ここで、結果を視覚化します。
このリポジトリが研究に役立つと思われる場合は、次の論文を引用することを検討してください。
@article { vangansbeke2024ldmseg ,
title = { a simple latent diffusion approach for panoptic segmentation and mask inpainting } ,
author = { Van Gansbeke, Wouter and De Brabandere, Bert } ,
journal = { arxiv preprint arxiv:2401.10227 } ,
year = { 2024 }
}ご不明な点がございましたら、主著者までご連絡ください。
このソフトウェアは、個人および研究での使用のみを許可するクリエイティブ コモンズ ライセンスに基づいてリリースされています。商用ライセンスについては、著者にお問い合わせください。ここでライセンスの概要を確認できます。
すべてのパブリック リポジトリ (コード内の参照も参照)、特に detectron2 ライブラリと diffusers ライブラリに感謝します。
