Skip to content

README.md

Latest commit

 

History

History
281 lines (226 loc) · 9.27 KB

README.md

File metadata and controls

281 lines (226 loc) · 9.27 KB
Raw

DOEI: Dual Optimization of Embedding Information for Attention-Enhanced Class Activation Maps

This is the code repository for the paper:

DOEI: Dual Optimization of Embedding Information for Attention-Enhanced Class Activation Maps

Hongjie Zhu, Zeyu Zhang†, Guansong Pang, Xu Wang, Shimin Wen, Yu Bai, Daji Ergu, Ying Cai* and Yang Zhao

*Corresponding author. †Project lead.

Neurocomputing 2026

[arXiv] [Model] [Paper with Code]

DOEI overview

Introduction

Weakly supervised semantic segmentation (WSSS) typically uses limited semantic annotations to obtain initial class activation maps (CAMs). DOEI optimizes embedding information through semantic-aware attention weight matrices: high-confidence tokens are amplified, low-confidence tokens are suppressed, and a hybrid-feature alignment module combines RGB values, embedding-guided features, and self-attention weights. This improves target feature propagation and decoupling for visual-transformer WSSS models.

This repository contains a PyTorch implementation of the DeiT/MCTformer-style weakly supervised pipeline used by DOEI. The main model performs multi-label image classification and can export class activation maps (CAMs). The psa/ and seg/ directories provide optional affinity propagation and segmentation stages for converting CAMs into semantic segmentation masks.

Repository layout

main.py                 classification training, evaluation, and CAM export
datasets.py             VOC12/COCO dataset loaders
models.py               MCTformer model definition
engine.py               training, mAP evaluation, and CAM generation
evaluation.py           CAM/segmentation mIoU evaluation
psa/                    pixel-affinity training and inference
seg/                    pseudo-label segmentation training and inference
voc12/                  image lists and VOC12 multi-label annotations

Installation

Use Python 3.8 or newer and a PyTorch build compatible with your CUDA driver. The code uses the older timm model registry API, so timm==0.4.12 is the recommended version.

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121
python -m pip install timm==0.4.12 numpy pillow pandas scikit-learn opencv-python dill scipy scikit-image tqdm

Install the CUDA-enabled PyTorch command that matches your machine if CUDA 12.1 is not appropriate. DenseCRF is only needed for the optional CRF paths:

python -m pip install pydensecrf

Run all commands from the repository root so that the local model and voc12/cls_labels.npy files are found correctly.

Data preparation

Download Pascal VOC 2012 and, if you use the augmented classification list, the VOC12 extra training images. Set VOC_ROOT to the directory containing the standard VOC layout:

$VOC_ROOT/
  JPEGImages/*.jpg
  Annotations/*.xml
  SegmentationClass/*.png
  ImageSets/Segmentation/*.txt

The repository already includes the image lists and multi-label annotations:

voc12/train_aug_id.txt
voc12/train_id.txt
voc12/val_id.txt
voc12/cls_labels.npy

The default VOC12 loader reads images from $VOC_ROOT/JPEGImages and uses voc12/train_aug_id.txt for training and voc12/val_id.txt for validation.

Model weights

GitHub is used for source code and small metadata only. The checkpoint files larger than GitHub's recommended file size are stored with their original paths in AIGeeksGroup/DOEI on Hugging Face:

File Purpose
deit_small_patch16_224-cd65a155.pth DeiT-S initialization used by models.py
weights/res38_cls.pth ResNet-38 classification initialization for PSA
saved_model/checkpoint.pth saved MCTformer checkpoint
saved_model_21_num_class/checkpoint.pth saved 21-class checkpoint
results/deit_small_MCTformerPlus/checkpoint.pth MCTformerPlus result checkpoint

Download them without putting credentials in the repository:

python -m pip install -U huggingface_hub
python - <<'PY'
from huggingface_hub import hf_hub_download

files = [
    "deit_small_patch16_224-cd65a155.pth",
    "weights/res38_cls.pth",
    "saved_model/checkpoint.pth",
    "saved_model_21_num_class/checkpoint.pth",
    "results/deit_small_MCTformerPlus/checkpoint.pth",
]
for name in files:
    hf_hub_download(
        repo_id="AIGeeksGroup/DOEI",
        filename=name,
        local_dir=".",
    )
PY

For a private Hugging Face clone, authenticate with huggingface-cli login or set HF_TOKEN in the shell; never commit that token.

Training

Multi-label classification

main.py trains deit_small_MCTformerPlus with the VOC12 multi-label objective and evaluates validation mAP after every epoch. The default run is 45 epochs, batch size 16, input size 224, and AdamW.

python main.py \
  --data-path "$VOC_ROOT" \
  --img-list voc12 \
  --model deit_small_MCTformerPlus \
  --output_dir output_models/doei \
  --epochs 45 \
  --batch-size 16 \
  --device cuda

Use --device cpu for a smoke test only; the default settings are intended for a CUDA GPU. A checkpoint and JSON-lines training log are written to --output_dir/checkpoint.pth and --output_dir/log.txt.

To fine-tune or evaluate a saved checkpoint, pass it through --finetune:

python main.py \
  --eval \
  --finetune output_models/doei/checkpoint.pth \
  --data-path "$VOC_ROOT" \
  --img-list voc12 \
  --device cuda

--eval reports classification mAP on the VOC12 validation list. In the current implementation --finetune is the checkpoint-loading option for this path; --resume is used by CAM generation below.

CAM generation

Generate multi-scale CAM .npy files and optional visualizations from a checkpoint trained by main.py:

python main.py \
  --gen_attention_maps \
  --resume output_models/doei/checkpoint.pth \
  --data-path "$VOC_ROOT" \
  --img-list voc12 \
  --attention-dir results/cam_png \
  --cam-npy-dir results/cam_npy \
  --device cuda

The generated .npy files store per-class CAMs. Their background threshold can be selected with evaluation.py against VOC segmentation masks.

Optional PSA and segmentation stages

The auxiliary scripts require pseudo-label/CRF directories produced by the preceding stage. Their complete argument lists are available with --help. Typical entry points are:

# Train the pixel-affinity network (the script imports psa/network).
cd psa
python train_aff.py \
  --weights ../weights/res38_cls.pth \
  --train_list voc12/train_aug.txt \
  --val_list voc12/val.txt \
  --voc12_root "$VOC_ROOT" \
  --la_crf_dir /path/to/low-affinity-crf \
  --ha_crf_dir /path/to/high-affinity-crf \
  --session_name ../results/resnet38_aff
cd ..

# Train a segmentation network from pseudo labels (the script imports seg/network).
cd seg
python train_seg.py \
  --network resnet38_seg \
  --init_weights /path/to/segmentation-init.pth \
  --list_path ../voc12/train_aug_id.txt \
  --img_path "$VOC_ROOT/JPEGImages" \
  --seg_pgt_path /path/to/pseudo-labels \
  --save_path ../results/segmentation
cd ..

Testing and segmentation evaluation

evaluation.py computes per-class IoU, mean IoU, false-positive rate, and false-negative rate from CAM predictions and VOC PNG masks. The list passed to --list must contain IDs for which both prediction and ground-truth files exist.

For CAM .npy predictions:

python evaluation.py \
  --list "$VOC_ROOT/ImageSets/Segmentation/val.txt" \
  --predict_dir results/cam_npy \
  --gt_dir "$VOC_ROOT/SegmentationClass" \
  --type npy \
  --start 0 \
  --end 100 \
  --logfile results/cam_eval.txt \
  --comment val

The script sweeps background thresholds from 0.00 to 0.99 and prints the best mIoU. To evaluate indexed PNG predictions instead, use --type png and point --predict_dir at the PNG directory.

For segmentation-network inference and mIoU reporting:

cd seg
python infer_seg.py \
  --weights /path/to/segmentation-checkpoint.pth \
  --network resnet38_seg \
  --gt_path "$VOC_ROOT/SegmentationClass" \
  --list_path ../voc12/val_id.txt \
  --img_path "$VOC_ROOT/JPEGImages" \
  --save_path ../results/seg_pred \
  --save_path_c ../results/seg_pred_palette \
  --scales 1.0 0.5 1.5
cd ..

The inference script writes per-class IoU and mIoU to results/seg_pred/result.txt.

Reproducibility notes

  • The default random seed is 0; set --seed to change it.
  • Classification uses 20 foreground VOC categories (nb_classes=20); the segmentation evaluator uses 21 labels including background.
  • Paths in the original scripts default to Windows examples. Always pass Linux paths explicitly on Linux.
  • GPU memory use depends on input size, batch size, and the selected number of workers. Reduce --batch-size and --num_workers when necessary.

License

No license file is currently included. Add the license required by your project or dataset terms before redistributing the repository.