ryh1218/occProhibited Items Segmentation via Occlusion-aware Bilayer Modeling (ICME 2025)
遮蔽認識バイレイヤーモデリングによる禁止物品セグメンテーションの公式 PyTorch 実装です。
コードは、Pytorch 2.1.2 と CUDA 11.8 を備えた Linux でテスト済みです。
必要に応じて、numpy<2.0.0 をインストールしてください（当方では 1.24.0 を使用しています）：
conda create -n occ python=3.10
conda activate occ
pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu118
pip install -U openmim
mim install mmcv==2.1.0
conda install mpi4py
pip install -U transformers==4.38.1 wandb==0.16.3 einops pycocotools shapely scipy terminaltables importlib peft==0.8.2 mat4py==0.6.0
conda install importlib-metadata
SAM の重み（例：SAM-vit-large）を事前にダウンロードし、'convert_model_weights.py' 内のパスを更新してください。重みは以下のスクリプトでダウンロードできます：
from transformers import SamModel
import os
hf_pretrain_name = "facebook/sam-vit-large"
cache_dir = f"{os.path.basename(hf_pretrain_name).replace('-', '_')}"
os.makedirs(cache_dir, exist_ok=True)
model = SamModel.from_pretrained(hf_pretrain_name, use_safetensors=False)
model.save_pretrained(cache_dir, safe_serialization=False)
その後、'convert_model_weights.py' スクリプトを実行して、Occluder Mask Decoder 用の初期重みを生成してください。
以下のリンクからデータセットをダウンロードしてください。PIXray は COCO 形式に変換済みであることに注意してください：
'convert_occlusion_annotation.py' を使用して、COCO 形式データセット用の遮蔽アノテーションを生成してください。独自の遮蔽アノテーション付きデータセットを生成するには、実行前にスクリプト内のパスを変更してください。
お使いの環境に合わせて、設定ファイル 'configs/occlusion/pid.py' と 'configs/occlusion/pix.py' のパスを更新してください。
当コードは以下でテスト済みです：
NVIDIA A800 Tensor Core GPU × 1
NVIDIA GeForce RTX 4090 GPU × 8
モデルをトレーニングするには：
# Single GPU, PIDray-A
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_train.sh configs/occlusion/pid.py 1
# Single GPU, PIXray-A
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_train.sh configs/occlusion/pix.py 1
# Multiple GPUs, PIDray-A
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 bash ./tools/dist_train.sh configs/occlusion/pid.py 8
# Multiple GPUs, PIXray-A
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 bash ./tools/dist_train.sh configs/occlusion/pix.py 8
モデルをテストするには、'CHECKPOINT_FILE' をチェックポイントのパスに設定し、'GPU_NUM' を GPU 構成に合わせて設定してください：
# PIDray-A
bash ./tools/dist_test.sh configs/occlusion/pid.py {CHECKPOINT_FILE} {GPU_NUM}
# PIXray-A
bash ./tools/dist_test.sh configs/occlusion/pix.py {CHECKPOINT_FILE} {GPU_NUM}
さらに、モデルのチェックポイントを以下のリンクで提供しています：Checkpoints
本実装は以下を基にしています：
これらの貢献に心より感謝いたします。
@inproceedings{ren2025prohibited,
title={Prohibited Items Segmentation via Occlusion-aware Bilayer Modeling},
author={Ren, Yunhan and Li, Ruihuang and Liu, Lingbo and Chen, Changwen},
booktitle={2025 IEEE International Conference on Multimedia and Expo (ICME)},
year={2025},
}