Segmentazione di Oggetti Proibiti tramite Modellazione a Doppio Strato Sensibile all'Occlusione (ICME 2025)

Implementazione PyTorch Ufficiale della Segmentazione di Oggetti Proibiti tramite Modellazione a Doppio Strato Sensibile all'Occlusione.

Link al paper: IEEE, Arxiv.

Per Iniziare

Dipendenze

Il codice è stato testato con successo su Linux con Pytorch 2.1.2 e CUDA 11.8.

Se necessario, installa numpy<2.0.0 (noi usiamo 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

Preparazione dei Pesi del Modello

Pre-scarica i pesi di SAM (ad es., SAM-vit-large) e aggiorna i percorsi in 'convert_model_weights.py'. I pesi possono essere scaricati usando il seguente 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)

Dopo aver fatto ciò, esegui lo script 'convert_model_weights.py' per generare i pesi iniziali per l'Occluder Mask Decoder.

Preparazione dei Dataset

Scarica i dataset dai seguenti link, nota che PIXray è stato convertito nel formato COCO:

PIDray-A

PIXray-A

Supporto per Dataset Personalizzati

Usa 'convert_occlusion_annotation.py' per generare annotazioni di occlusione per dataset in formato COCO. Per generare i tuoi dataset con annotazioni di occlusione, modifica i percorsi nello script prima dell'esecuzione.

Modifica dei File di Configurazione

Aggiorna i percorsi nei file di configurazione 'configs/occlusion/pid.py' e 'configs/occlusion/pix.py' in base alla tua configurazione.

Addestramento e Test

Il nostro codice è stato testato con successo su:

• 1 × GPU NVIDIA A800 Tensor Core

• 8 × GPU NVIDIA GeForce RTX 4090

Per addestrare il nostro modello:

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

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

Per testare il nostro modello, imposta 'CHECKPOINT_FILE' sul percorso del tuo checkpoint e configura 'GPU_NUM' in base alla tua configurazione 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}

Inoltre, forniamo i checkpoint del modello al seguente link: Checkpoints

Ringraziamenti

Questa implementazione si basa su:

• RSPrompter

• MMDetection

Apprezziamo sinceramente i loro contributi.

Cita questo paper

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