基于遮挡感知双层建模的违禁物品分割（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

训练我们的模型：

# 单 GPU，PIDray-A
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_train.sh configs/occlusion/pid.py 1
# 单 GPU，PIXray-A
CUDA_VISIBLE_DEVICES=0 bash ./tools/dist_train.sh configs/occlusion/pix.py 1

# 多 GPU，PIDray-A
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 bash ./tools/dist_train.sh configs/occlusion/pid.py 8
# 多 GPU，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},
}