Segmentation d'objets interdits via une modélisation bicouche sensible à l'occlusion (ICME 2025)

Implémentation officielle PyTorch de la segmentation d'objets interdits via une modélisation bicouche sensible à l'occlusion.

Lien vers l'article : IEEE, Arxiv.

Pour commencer

Dépendances

Le code a été testé avec succès sur Linux avec PyTorch 2.1.2 et CUDA 11.8.

Si nécessaire, installez numpy<2.0.0 (nous utilisons la version 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

Préparer les poids du modèle

Téléchargez à l'avance les poids SAM (par exemple, SAM-vit-large) et mettez à jour les chemins dans 'convert_model_weights.py'. Les poids peuvent être téléchargés à l'aide du script suivant :

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)

Après cela, exécutez le script 'convert_model_weights.py' pour générer les poids initiaux de l'Occluder Mask Decoder.

Préparer les jeux de données

Téléchargez les jeux de données à partir des liens suivants ; notez que PIXray a été converti au format COCO :

PIDray-A

PIXray-A

Prise en charge des jeux de données personnalisés

Utilisez 'convert_occlusion_annotation.py' pour générer des annotations d'occlusion pour les jeux de données au format COCO. Pour générer vos propres jeux de données avec annotations d'occlusion, modifiez les chemins dans le script avant de l'exécuter.

Modifier les fichiers de configuration

Mettez à jour les chemins dans les fichiers de configuration 'configs/occlusion/pid.py' et 'configs/occlusion/pix.py' selon votre configuration.

Entraînement et test

Notre code a été testé avec succès sur :

• 1 × NVIDIA A800 Tensor Core GPU

• 8 × NVIDIA GeForce RTX 4090 GPUs

Pour entraîner notre modèle :

# 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

Pour tester notre modèle, définissez 'CHECKPOINT_FILE' sur le chemin de votre point de contrôle, et configurez 'GPU_NUM' en fonction de votre configuration 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}

De plus, nous fournissons des points de contrôle du modèle via le lien suivant : Points de contrôle

Remerciements

Cette implémentation s'appuie sur :

• RSPrompter

• MMDetection

Nous apprécions sincèrement leurs contributions.

Citer cet article

@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},
}