Segmentación de artículos prohibidos mediante modelado bicapa consciente de la oclusión (ICME 2025)

Implementación oficial en PyTorch de la segmentación de artículos prohibidos mediante modelado bicapa consciente de la oclusión.

Enlace del artículo: IEEE, Arxiv.

Cómo empezar

Dependencias

El código se ha probado con éxito en Linux con Pytorch 2.1.2 y CUDA 11.8.

Si es necesario, instala numpy<2.0.0 (nosotros 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 los pesos del modelo

Descarga previamente los pesos de SAM (p. ej., SAM-vit-large) y actualiza las rutas en 'convert_model_weights.py'. Los pesos se pueden descargar con el siguiente 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)

Después de hacerlo, ejecuta el script 'convert_model_weights.py' para generar los pesos iniciales para el decodificador de máscaras de oclusor.

Preparar los conjuntos de datos

Descarga los conjuntos de datos desde los siguientes enlaces; ten en cuenta que PIXray se ha transformado al formato COCO:

PIDray-A

PIXray-A

Soporte para conjuntos de datos personalizados

Usa 'convert_occlusion_annotation.py' para generar anotaciones de oclusión para conjuntos de datos en formato COCO. Para generar tus propios conjuntos de datos con anotaciones de oclusión, modifica las rutas en el script antes de ejecutarlo.

Modificar los archivos de configuración

Actualiza las rutas en los archivos de configuración 'configs/occlusion/pid.py' y 'configs/occlusion/pix.py' según tu configuración.

Entrenar y probar

Nuestro código se ha probado con éxito en:

• 1 × GPU NVIDIA A800 Tensor Core

• 8 × GPUs NVIDIA GeForce RTX 4090

Para entrenar nuestro modelo:

# 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

Para probar nuestro modelo, establece 'CHECKPOINT_FILE' en la ruta de tu checkpoint y configura 'GPU_NUM' para que coincida con tu configuración de 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}

Además, proporcionamos checkpoints del modelo en el siguiente enlace: Checkpoints

Agradecimientos

Esta implementación se basa en:

• RSPrompter

• MMDetection

Agradecemos sinceramente sus contribuciones.

Citar este artículo

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