app_infer.py

import sys
import os

# ================= 关键修改开始 =================
# 把 defake 目录加入到系统路径，解决 ModuleNotFoundError
current_dir = os.path.dirname(os.path.abspath(__file__))
defake_dir = os.path.join(current_dir, "defake")
sys.path.append(defake_dir)
# ================= 关键修改结束 =================

from PIL import Image
import tempfile
from defake.test_api import load_models, predict_image

# 全局只加载一次模型
_models = None

def get_models():
    global _models
    if _models is None:
        _models = load_models()
    return _models

def run_infer_from_image(pil_image: Image.Image):
    """
    输入：PIL Image
    输出：dict，包含 fake_score、is_fake 等
    """
    models = get_models()

    # 暂存到临时文件
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp:
        pil_image.save(tmp.name)
        pred, probs = predict_image(tmp.name, models)

    # 假设类别 1 是 fake，0 是 real
    fake_score = float(probs[1])
    is_fake = (pred == 1)

    return {
        "fake_score": fake_score,
        "is_fake": is_fake,
        "probs": probs,
        "pred_class": int(pred),
    }