retofan23333/UniDG-Defect-Recog-Reward-Release

UniDG-Defect-Recog-Reward

Defect-Recog-Reward (final_model.pth, SegFormer-B5) for UniDG (UniDG: Unified Defect Generation via Defect-Context Editing, NeurIPS 2026).

[Paper] [Code] [UniDG-SFT-LoRA] [UniDG-RFT-LoRA]

Overview

This checkpoint is the unified 28-class (background + 27 defect categories) SegFormer-B5 (nvidia/mit-b5) recognition and segmentation reward model used in Consistency-RFT. It scores each generated defect image by the average of classification match, pixel AUROC, pixel AP, and PRO.

Usage in UniDG Consistency-RFT

hf download retofan23333/UniDG-Defect-Recog-Reward-Release --local-dir checkpoints/UniDG-Defect-Recog-Reward-Release

Point recog_model in train/rft/config.yaml to checkpoints/UniDG-Defect-Recog-Reward-Release/final_model.pth.

Citation

@inproceedings{fan2026unidg,
  title={UniDG: Unified Defect Generation via Defect-Context Editing},
  author={Fan, Yuanting and Liu, Jun and Gao, Bin-Bin and Chen, Xiaochen and Lin, Yuhuan and Dai, Zhewei and Zhan, Jiawei and Wang, Chengjie},
  booktitle={NeurIPS},
  volume={39},
  year={2026},
  publisher={Curran Associates, Inc.}
}