#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""data_audit.py —— 图像数据体检（7.5 数据管线的学生版工具）
对一个图像目录做：损坏检测、EXIF 旋转检查、尺寸统计、dHash 模糊去重，
输出 JSON 报告 + 可疑重复对清单。不修改任何原始文件。
依赖：pip install pillow numpy
用法：
  python data_audit.py --dir data/raw --report audit.json [--hash-dist 6]
"""
import argparse, json, os, sys
import numpy as np
from PIL import Image, ImageOps

def dhash(img, size=8):
    """差异哈希：缩到 (size+1, size) 后按水平梯度生成 size*size 位。纯 PIL/numpy，无额外依赖。"""
    g = np.asarray(ImageOps.grayscale(img).resize((size + 1, size), Image.LANCZOS), dtype=np.int16)
    return (g[:, 1:] > g[:, :-1]).flatten()

def main():
    ap = argparse.ArgumentParser(description="图像数据体检")
    ap.add_argument("--dir", required=True, help="图像目录（递归）")
    ap.add_argument("--report", default="audit_report.json", help="输出报告路径")
    ap.add_argument("--hash-dist", type=int, default=6, help="dHash 汉明距离 <= 该值视为可疑重复")
    args = ap.parse_args()

    exts = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
    files = [os.path.join(r, f) for r, _, fs in os.walk(args.dir) for f in fs
             if os.path.splitext(f)[1].lower() in exts]
    print(f"发现 {len(files)} 个图像文件，开始体检……")

    corrupt, exif_rotated, hashes, sizes, ars = [], [], [], [], []
    for i, p in enumerate(files):
        if i % 200 == 0:
            print(f"  {i}/{len(files)}")
        try:
            with Image.open(p) as im:
                im.verify()                                   # 第一层：文件完整性
            with Image.open(p) as im:
                exif = im.getexif()
                if exif and exif.get(274, 1) != 1:            # 274 = Orientation 标签
                    exif_rotated.append(p)
                im = ImageOps.exif_transpose(im)              # 统一转正后再取特征
                w, h = im.size
                sizes.append((w, h)); ars.append(w / max(h, 1))
                hashes.append((p, dhash(im)))
        except Exception as e:
            corrupt.append({"file": p, "error": str(e)})

    # 模糊去重：暴力两两比较（<5 万张可接受；更大规模请上 LSH）
    W = args.hash_dist
    dup_pairs = []
    n = len(hashes)
    for i in range(n):
        pi, hi = hashes[i]
        for j in range(i + 1, n):
            if (hi != hashes[j][1]).sum() <= W:
                dup_pairs.append([pi, hashes[j][0]])

    ws = [s[0] for s in sizes]; hs = [s[1] for s in sizes]
    report = {
        "dir": args.dir,
        "total": len(files),
        "corrupt": corrupt,
        "exif_rotated_count": len(exif_rotated),
        "exif_rotated_sample": exif_rotated[:10],
        "size_stats": {
            "width_min_max": [min(ws), max(ws)] if ws else None,
            "height_min_max": [min(hs), max(hs)] if hs else None,
            "aspect_ratio_p5_p95": [float(np.percentile(ars, 5)), float(np.percentile(ars, 95))] if ars else None,
        },
        "suspected_dup_pairs": dup_pairs,
        "dup_pair_count": len(dup_pairs),
        "hash_dist_threshold": W,
    }
    with open(args.report, "w", encoding="utf-8") as f:
        json.dump(report, f, ensure_ascii=False, indent=2)

    print("=" * 56)
    print(f"总数 {len(files)} | 损坏 {len(corrupt)} | EXIF 需转正 {len(exif_rotated)} | 可疑重复对 {len(dup_pairs)}")
    print(f"长宽比 p5/p95 = {report['size_stats']['aspect_ratio_p5_p95']}（差异大 → 训练前分桶或统一裁剪）")
    print(f"完整报告 → {args.report}")
    print("提示：EXIF 转正样本请在入库管线统一 ImageOps.exif_transpose（7.5 §2）；重复对人工抽查后再决定剔除。")

if __name__ == "__main__":
    main()