ReFor

PRNU 기반 이미지 위·변조 탐지 및 생성 모델 소속 분류를 위한 PyTorch 프로젝트입니다.
2단계 파이프라인 구조로, 먼저 Real/Fake 이진 분류를 수행하고 Fake로 판별된 경우 생성 모델 Attribution 분류를 수행합니다.

---

파일 구조

ReFor/
├── scripts/
│   ├── main.py
│   └── bootstrap.py
├── refor/
│   ├── core/
│   │   └── env.py
│   ├── data/
│   │   ├── data.py
│   │   └── residual.py
│   ├── eval/
│   │   ├── e2e_eval.py
│   │   └── eval_balanced.py
│   ├── models/
│   │   └── models.py
│   ├── train/
│   │   └── train_utils.py
│   └── infer/
│       └── inference.py
├── ReFor.ipynb
├── ReForwithUnknownTest.ipynb
└── run.sh

---

데이터셋 준비

데이터셋은 기본적으로 다음 경로 구조를 가정합니다:

/content/prnu_dataset/data/
├── ModelA/
│   ├── 0_real/   # 진짜 이미지
│   └── 1_fake/   # ModelA가 생성한 이미지
├── ModelB/
│   ├── 0_real/
│   └── 1_fake/
└── ...

• ModelA, ModelB 등 폴더 이름은 생성 모델 이름이어야 합니다.
• 각 모델 폴더 안에는 반드시 0_real/, 1_fake/ 하위 폴더가 있어야 합니다.
• 최소 하나 이상의 모델 폴더가 필요합니다.
• 실험에서는 DIF dataset을 사용했습니다.

---

실행 방법

Colab / 셸에서:

# 1) 의존성 설치
python ReFor/scripts/bootstrap.py

# 2) 환경 초기화 확인 (옵션)
python -m refor.core.env

# 3) 학습 및 평가
python ReFor/scripts/main.py

---

추론 예시

단일 이미지에 대해 2단계 추론을 수행하려면

from refor.inference.inference import predict_two_stage
from refor.models.models import ResNetHead
from refor.data.data import make_attr_loaders
import torch

DATA_ROOT = "/content/prnu_dataset/data"
_, _, attr_classes = make_attr_loaders(DATA_ROOT)

binary = ResNetHead(2)
binary.load_state_dict(torch.load(f"{DATA_ROOT}/_ckpts_multi/binary_full_best.pth"))

attr = ResNetHead(len(attr_classes))
attr.load_state_dict(torch.load(f"{DATA_ROOT}/_ckpts_multi/attr_full_best.pth"))

result = predict_two_stage("test.jpg", binary, attr, attr_classes, tta=True, fake_threshold=0.5)
print(result)

⚠️ 수행하려면 실행 방법을 따라 실행 후 .pth 파일을 생성한 후 해당 모델 파일을 사용해야 합니다.

---

참고

• GPU 환경에서 실행하는 것을 권장합니다.