Occlusion-aware Bilayer Modeling के माध्यम से निषिद्ध वस्तुओं का विभाजन (ICME 2025)

Occlusion-aware Bilayer Modeling के माध्यम से निषिद्ध वस्तुओं के विभाजन का आधिकारिक 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)

ऐसा करने के बाद, Occluder Mask Decoder के लिए प्रारंभिक वेट उत्पन्न करने हेतु 'convert_model_weights.py' स्क्रिप्ट चलाएँ।

डेटासेट तैयार करें

निम्नलिखित लिंक से डेटासेट डाउनलोड करें, ध्यान दें कि PIXray को COCO-प्रारूप में बदल दिया गया है:

PIDray-A

PIXray-A

कस्टम डेटासेट समर्थन

'convert_occlusion_annotation.py' का उपयोग करके COCO-प्रारूप डेटासेट के लिए occlusion एनोटेशन उत्पन्न करें। अपने स्वयं के occlusion-एनोटेटेड डेटासेट उत्पन्न करने के लिए, निष्पादन से पहले स्क्रिप्ट में पथ संशोधित करें।

कॉन्फ़िग फ़ाइलें संशोधित करें

अपनी सेटअप के अनुसार कॉन्फ़िगरेशन फ़ाइलों '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' को अपने checkpoint पथ पर सेट करें, और '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},
}