Segmentação de Itens Proibidos via Modelagem Bilayer Ciente de Oclusão (ICME 2025)

Implementação Oficial em PyTorch da Segmentação de Itens Proibidos via Modelagem Bilayer Ciente de Oclusão.

Link do artigo: IEEE, Arxiv.

Começando

Dependências

O código foi testado com sucesso em Linux com Pytorch 2.1.2 e CUDA 11.8.

Se necessário, instale numpy<2.0.0 (nós usamos 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

Preparar Pesos do Modelo

Baixe previamente os pesos do SAM (ex.: SAM-vit-large) e atualize os caminhos em 'convert_model_weights.py'. Os pesos podem ser baixados usando o seguinte script:

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)

Depois disso, execute o script 'convert_model_weights.py' para gerar os pesos iniciais para o Occluder Mask Decoder.

Preparar Conjuntos de Dados

Baixe os conjuntos de dados dos seguintes links. Observe que o PIXray foi convertido para o formato COCO:

PIDray-A

PIXray-A

Suporte a Conjuntos de Dados Personalizados

Use 'convert_occlusion_annotation.py' para gerar anotações de oclusão para conjuntos de dados no formato COCO. Para gerar seus próprios conjuntos de dados com anotações de oclusão, modifique os caminhos no script antes da execução.

Modificar Arquivos de Configuração

Atualize os caminhos nos arquivos de configuração 'configs/occlusion/pid.py' e 'configs/occlusion/pix.py' de acordo com seu ambiente.

Treinar e Testar

Nosso código foi testado com sucesso em:

• 1 × GPU NVIDIA A800 Tensor Core

• 8 × GPUs NVIDIA GeForce RTX 4090

Para treinar nosso modelo:

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

# Múltiplas GPUs, PIDray-A
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 bash ./tools/dist_train.sh configs/occlusion/pid.py 8
# Múltiplas GPUs, PIXray-A
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 bash ./tools/dist_train.sh configs/occlusion/pix.py 8

Para testar nosso modelo, defina o 'CHECKPOINT_FILE' para o caminho do seu checkpoint e configure o 'GPU_NUM' para corresponder ao seu ambiente de GPUs:

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

Além disso, fornecemos checkpoints do modelo no seguinte link: Checkpoints

Agradecimentos

Esta implementação é baseada em:

• RSPrompter

• MMDetection

Agradecemos sinceramente suas contribuições.

Citar este artigo

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