Сегментация запрещённых предметов с помощью двухслойного моделирования с учётом окклюзии (ICME 2025)

Официальная реализация на PyTorch метода сегментации запрещённых предметов с помощью двухслойного моделирования с учётом окклюзии.

Ссылка на статью: IEEE, Arxiv.

Начало работы

Зависимости

Код успешно протестирован на Linux с Pytorch 2.1.2 и CUDA 11.8.

При необходимости установите 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.

PIDray-A

PIXray-A

Поддержка пользовательских наборов данных

Используйте 'convert_occlusion_annotation.py' для генерации аннотаций окклюзии для наборов данных в формате COCO. Чтобы создать собственные наборы данных с аннотациями окклюзии, измените пути в скрипте перед запуском.

Изменение файлов конфигурации

Обновите пути в файлах конфигурации 'configs/occlusion/pid.py' и 'configs/occlusion/pix.py' в соответствии с вашей конфигурацией.

Обучение и тестирование

Наш код был успешно протестирован на:

• 1 × NVIDIA A800 Tensor Core GPU
• 8 × NVIDIA GeForce RTX 4090 GPU

Для обучения нашей модели:

# 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}

Кроме того, мы предоставляем чекпоинты модели по следующей ссылке: Чекпоинты

Благодарности

Эта реализация основана на:

• RSPrompter

• MMDetection

Мы искренне ценим их вклад.

Как цитировать эту статью

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