TraceDetect-AI / tune_thresholds.py
nchdlhbctm's picture
Upload 13 files
646535a verified
Raw
History Blame Contribute Delete
1.87 kB
import os
import cv2
import numpy as np
from image_module import get_lbp_entropy, get_dct_high_freq_energy, get_fft_symmetry_error
def process_folder(folder_path, label):
results = []
print(f"正在分析 [{label}] 样本集: {folder_path}")
for filename in os.listdir(folder_path):
if filename.lower().endswith(('.png', '.jpg', '.jpeg')):
filepath = os.path.join(folder_path, filename)
# 读取图片并转换为灰度图
img = cv2.imdecode(np.fromfile(filepath, dtype=np.uint8), cv2.IMREAD_COLOR)
if img is None: continue
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# 提取特征
lbp = get_lbp_entropy(gray)
dct = get_dct_high_freq_energy(gray)
fft = get_fft_symmetry_error(gray)
results.append((lbp, dct, fft))
if not results:
return None
results_arr = np.array(results)
print(f"--- {label} 统计结果 ({len(results)}张) ---")
print(
f"LBP 熵 : 平均 {np.mean(results_arr[:, 0]):.3f} | 范围 [{np.min(results_arr[:, 0]):.3f} - {np.max(results_arr[:, 0]):.3f}]")
print(
f"DCT 占比: 平均 {np.mean(results_arr[:, 1]):.3f} | 范围 [{np.min(results_arr[:, 1]):.3f} - {np.max(results_arr[:, 1]):.3f}]")
print(
f"FFT 误差: 平均 {np.mean(results_arr[:, 2]):.3f} | 范围 [{np.min(results_arr[:, 2]):.3f} - {np.max(results_arr[:, 2]):.3f}]\n")
return results_arr
if __name__ == "__main__":
# 请在当前目录下新建这两个文件夹,分别放几十张真实的和AI生成的图进去
real_folder = "./data/real"
ai_folder = "./data/ai"
print("开始执行特征分布寻优...\n")
process_folder(real_folder, "真实图像")
process_folder(ai_folder, "AI生成图像")