폐색 인지 이중 레이어 모델링을 통한 금지 품목 분할 (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 GPUs

모델을 학습하려면:

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

또한 다음 링크에서 모델 체크포인트를 제공합니다: 체크포인트

감사의 말

이 구현은 다음을 기반으로 합니다:

그들의 기여에 진심으로 감사드립니다.

논문 인용

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

Occ

인기