training/evaluate_image_detector.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path
from typing import Any, Dict, List

import numpy as np
from PIL import Image


IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
REPO_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO_ROOT / "backend"))

from analyzers.image_decision import AI_DETECTOR_MIN, AUTHENTIC_DETECTOR_MAX  # noqa: E402


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate a TruthShield image model on an untouched split.")
    parser.add_argument("--model-dir", default="training/models/truthshield-image-detector-v2")
    parser.add_argument("--data-dir", default="training/data/defactify_sample")
    parser.add_argument("--split", default="test")
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--max-samples", type=int, default=0)
    parser.add_argument(
        "--max-per-class",
        type=int,
        default=0,
        help="Deterministically sample at most this many images from each class (preferred for balanced reports).",
    )
    parser.add_argument("--max-misclassified", type=int, default=40)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", default="")
    parser.add_argument("--audit-report", default="training/data/image_dataset_audit.json")
    parser.add_argument("--include-audit-leakage", action="store_true")
    return parser.parse_args()


def main() -> None:
    import torch
    from transformers import AutoImageProcessor, AutoModelForImageClassification

    args = parse_args()
    model_dir = Path(args.model_dir)
    split_dir = Path(args.data_dir) / args.split
    if not model_dir.exists():
        raise SystemExit(f"Model folder not found: {model_dir}")
    if not split_dir.exists():
        raise SystemExit(f"Dataset split not found: {split_dir}")

    processor = AutoImageProcessor.from_pretrained(str(model_dir))
    model = AutoModelForImageClassification.from_pretrained(str(model_dir))
    model.eval()
    label2id = {str(label): int(index) for label, index in model.config.label2id.items()}
    id2label = {int(index): str(label) for index, label in model.config.id2label.items()}
    samples = _collect_samples(split_dir, label2id)
    excluded_paths: List[str] = []
    if not args.include_audit_leakage:
        samples, excluded_paths = _exclude_audit_leakage(
            samples,
            data_dir=Path(args.data_dir),
            split=args.split,
            audit_path=Path(args.audit_report),
        )
    if args.max_per_class > 0:
        samples = _sample_per_class(samples, max_per_class=args.max_per_class, seed=args.seed)
    if args.max_samples > 0 and len(samples) > args.max_samples:
        rng = np.random.default_rng(args.seed)
        chosen = sorted(rng.choice(len(samples), size=args.max_samples, replace=False).tolist())
        samples = [samples[index] for index in chosen]
    if not samples:
        raise SystemExit("No supported images were found in the requested split.")

    true_labels: List[int] = []
    predicted_labels: List[int] = []
    probabilities: List[List[float]] = []
    misclassified: List[Dict[str, Any]] = []
    print(f"Evaluating {len(samples)} images from {split_dir}...", flush=True)
    with torch.inference_mode():
        for start in range(0, len(samples), max(1, args.batch_size)):
            batch = samples[start : start + max(1, args.batch_size)]
            images = []
            for path, _ in batch:
                with Image.open(path) as image:
                    images.append(image.convert("RGB"))
            encoded = processor(images=images, return_tensors="pt")
            logits = model(**encoded).logits
            batch_probabilities = torch.softmax(logits, dim=-1).cpu().numpy()
            batch_predictions = np.argmax(batch_probabilities, axis=1)
            for (path, true_label), prediction, scores in zip(batch, batch_predictions, batch_probabilities):
                true_labels.append(true_label)
                predicted_labels.append(int(prediction))
                probabilities.append([float(value) for value in scores])
                if int(prediction) != true_label and len(misclassified) < max(0, args.max_misclassified):
                    misclassified.append(
                        {
                            "path": str(path),
                            "expected": id2label[true_label],
                            "predicted": id2label[int(prediction)],
                            "confidence": round(float(scores[int(prediction)]), 6),
                            "scores": {
                                id2label[index]: round(float(score), 6)
                                for index, score in enumerate(scores)
                            },
                        }
                    )
            completed = min(start + len(batch), len(samples))
            if completed % 500 == 0 or completed == len(samples):
                print(f"  {completed}/{len(samples)}", flush=True)

    report = build_report(
        np.asarray(true_labels, dtype=np.int64),
        np.asarray(predicted_labels, dtype=np.int64),
        np.asarray(probabilities, dtype=np.float64),
        id2label,
    )
    report.update(
        {
            "model_dir": str(model_dir.resolve()),
            "data_split": str(split_dir.resolve()),
            "sample_count": len(samples),
            "audit_leakage_excluded_count": len(excluded_paths),
            "audit_leakage_excluded_paths": excluded_paths,
            "misclassified_examples": misclassified,
            "sampling": {
                "seed": args.seed,
                "max_samples": args.max_samples,
                "max_per_class": args.max_per_class,
            },
        }
    )
    output_path = Path(args.output) if args.output else model_dir / f"truthshield_{args.split}_report.json"
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
    print(
        json.dumps(
            {
                key: report[key]
                for key in (
                    "accuracy",
                    "balanced_accuracy",
                    "macro_f1",
                    "legacy_frontend_likely_ai_decision",
                    "three_way_detection",
                )
            },
            indent=2,
        )
    )
    print(f"Full report: {output_path.resolve()}", flush=True)


def _collect_samples(split_dir: Path, label2id: Dict[str, int]) -> List[tuple[Path, int]]:
    samples: List[tuple[Path, int]] = []
    for label, label_id in sorted(label2id.items(), key=lambda item: item[1]):
        folder = split_dir / label
        if not folder.exists():
            print(f"Warning: missing label folder {folder}", flush=True)
            continue
        for path in sorted(folder.rglob("*")):
            if path.is_file() and path.suffix.lower() in IMAGE_EXTENSIONS:
                samples.append((path, label_id))
    return samples


def _sample_per_class(
    samples: List[tuple[Path, int]],
    max_per_class: int,
    seed: int,
) -> List[tuple[Path, int]]:
    rng = np.random.default_rng(seed)
    selected: List[tuple[Path, int]] = []
    labels = sorted({label for _, label in samples})
    for label in labels:
        members = [sample for sample in samples if sample[1] == label]
        if len(members) > max_per_class:
            indices = sorted(rng.choice(len(members), size=max_per_class, replace=False).tolist())
            members = [members[index] for index in indices]
        selected.extend(members)
    return selected


def _exclude_audit_leakage(
    samples: List[tuple[Path, int]],
    data_dir: Path,
    split: str,
    audit_path: Path,
) -> tuple[List[tuple[Path, int]], List[str]]:
    if not audit_path.exists():
        print(f"Warning: audit report not found; no leakage exclusions applied: {audit_path}", flush=True)
        return samples, []
    audit = json.loads(audit_path.read_text(encoding="utf-8"))
    excluded_relative = set()
    group_keys = (
        "cross_split_exact_duplicates",
        "cross_split_normalized_pixel_duplicates",
        "cross_split_perceptual_hash_collisions",
        "cross_split_perceptual_near_matches",
    )
    for key in group_keys:
        for group in audit.get(key, []):
            paths = [str(path).replace("/", "\\") for path in group.get("paths", [])]
            if len({_path_split(path) for path in paths}) <= 1:
                continue
            for path in paths:
                if _path_split(path) == split:
                    excluded_relative.add(path.lower())
    kept = []
    excluded = []
    for path, label in samples:
        relative = str(path.relative_to(data_dir)).replace("/", "\\")
        if relative.lower() in excluded_relative:
            excluded.append(relative)
        else:
            kept.append((path, label))
    if excluded:
        print(f"Excluded {len(excluded)} {split} images flagged as cross-split perceptual leakage.", flush=True)
    return kept, sorted(excluded)


def _path_split(path: str) -> str:
    normalized = path.replace("/", "\\")
    return normalized.split("\\", 1)[0] if normalized else "unknown"


def build_report(
    true_labels: np.ndarray,
    predictions: np.ndarray,
    probabilities: np.ndarray,
    id2label: Dict[int, str],
) -> Dict[str, Any]:
    label_ids = sorted(id2label)
    confusion = np.zeros((len(label_ids), len(label_ids)), dtype=np.int64)
    for expected, predicted in zip(true_labels, predictions):
        confusion[int(expected), int(predicted)] += 1

    classes: Dict[str, Any] = {}
    recalls = []
    f1_scores = []
    for label_id in label_ids:
        true_positive = int(confusion[label_id, label_id])
        false_positive = int(confusion[:, label_id].sum() - true_positive)
        false_negative = int(confusion[label_id, :].sum() - true_positive)
        support = int(confusion[label_id, :].sum())
        precision = true_positive / max(1, true_positive + false_positive)
        recall = true_positive / max(1, true_positive + false_negative)
        f1 = 2.0 * precision * recall / max(1e-12, precision + recall)
        recalls.append(recall)
        f1_scores.append(f1)
        classes[id2label[label_id]] = {
            "precision": round(precision, 6),
            "recall": round(recall, 6),
            "f1": round(f1, 6),
            "support": support,
        }

    ai_label_id = next((index for index, label in id2label.items() if label == "ai_generated"), None)
    binary: Dict[str, Any] = {"available": False}
    if ai_label_id is not None:
        ai_truth = true_labels == ai_label_id
        ai_scores = probabilities[:, ai_label_id]
        ai_predictions = ai_scores >= 0.5
        true_positive = int(np.sum(ai_predictions & ai_truth))
        false_positive = int(np.sum(ai_predictions & ~ai_truth))
        true_negative = int(np.sum(~ai_predictions & ~ai_truth))
        false_negative = int(np.sum(~ai_predictions & ai_truth))
        binary = {
            "available": True,
            "threshold": 0.5,
            "accuracy": round((true_positive + true_negative) / max(1, len(ai_truth)), 6),
            "precision": round(true_positive / max(1, true_positive + false_positive), 6),
            "recall": round(true_positive / max(1, true_positive + false_negative), 6),
            "specificity": round(true_negative / max(1, true_negative + false_positive), 6),
            "false_positive_rate": round(false_positive / max(1, true_negative + false_positive), 6),
            "false_negative_rate": round(false_negative / max(1, true_positive + false_negative), 6),
            "roc_auc": round(_binary_auc(ai_truth.astype(np.int64), ai_scores), 6),
            "expected_calibration_error": round(_expected_calibration_error(ai_truth, ai_scores), 6),
            "confusion_matrix": {
                "true_ai": true_positive,
                "false_ai_alarm": false_positive,
                "true_non_ai": true_negative,
                "missed_ai": false_negative,
            },
        }
        legacy = _likely_ai_threshold_metrics(ai_truth, ai_scores, threshold=0.70)
        three_way = _three_way_metrics(
            ai_truth,
            ai_scores,
            authentic_max=AUTHENTIC_DETECTOR_MAX,
            ai_min=AI_DETECTOR_MIN,
        )
        threshold_analysis = [
            _likely_ai_threshold_metrics(ai_truth, ai_scores, threshold=threshold)
            for threshold in (0.50, 0.70, 0.80, 0.85, 0.90, 0.95)
        ]
    else:
        legacy = {"available": False}
        three_way = {"available": False}
        threshold_analysis = []

    return {
        "accuracy": round(float(np.mean(predictions == true_labels)), 6),
        "balanced_accuracy": round(float(np.mean(recalls)), 6),
        "macro_f1": round(float(np.mean(f1_scores)), 6),
        "classes": classes,
        "labels": [id2label[index] for index in label_ids],
        "confusion_matrix_rows_expected_columns_predicted": confusion.tolist(),
        "binary_ai_detection": binary,
        "legacy_frontend_likely_ai_decision": legacy,
        "three_way_detection": three_way,
        "likely_ai_threshold_analysis": threshold_analysis,
    }


def _likely_ai_threshold_metrics(ai_truth: np.ndarray, scores: np.ndarray, threshold: float) -> Dict[str, Any]:
    predicted_ai = scores >= threshold
    true_ai = int(np.sum(predicted_ai & ai_truth))
    false_ai = int(np.sum(predicted_ai & ~ai_truth))
    ai_count = int(np.sum(ai_truth))
    non_ai_count = int(np.sum(~ai_truth))
    return {
        "available": True,
        "threshold": threshold,
        "true_ai": true_ai,
        "false_ai_alarm": false_ai,
        "ai_count": ai_count,
        "non_ai_count": non_ai_count,
        "ai_recall": round(true_ai / max(1, ai_count), 6),
        "false_positive_rate": round(false_ai / max(1, non_ai_count), 6),
        "precision": round(true_ai / max(1, true_ai + false_ai), 6),
    }


def _three_way_metrics(
    ai_truth: np.ndarray,
    scores: np.ndarray,
    authentic_max: float,
    ai_min: float,
) -> Dict[str, Any]:
    likely_ai = scores >= ai_min
    likely_authentic = scores <= authentic_max
    inconclusive = ~(likely_ai | likely_authentic)
    ai_count = int(np.sum(ai_truth))
    non_ai_count = int(np.sum(~ai_truth))
    true_ai = int(np.sum(likely_ai & ai_truth))
    false_ai = int(np.sum(likely_ai & ~ai_truth))
    true_authentic = int(np.sum(likely_authentic & ~ai_truth))
    ai_mislabeled_authentic = int(np.sum(likely_authentic & ai_truth))
    decisive = int(np.sum(likely_ai | likely_authentic))
    decisive_correct = true_ai + true_authentic
    return {
        "available": True,
        "authentic_max": authentic_max,
        "ai_min": ai_min,
        "counts": {
            "likely_ai": int(np.sum(likely_ai)),
            "likely_authentic": int(np.sum(likely_authentic)),
            "inconclusive": int(np.sum(inconclusive)),
            "true_ai": true_ai,
            "false_ai_alarm": false_ai,
            "true_authentic": true_authentic,
            "ai_mislabeled_authentic": ai_mislabeled_authentic,
            "ai_inconclusive": int(np.sum(inconclusive & ai_truth)),
            "non_ai_inconclusive": int(np.sum(inconclusive & ~ai_truth)),
        },
        "false_positive_rate": round(false_ai / max(1, non_ai_count), 6),
        "false_negative_rate_as_authentic": round(ai_mislabeled_authentic / max(1, ai_count), 6),
        "ai_recall": round(true_ai / max(1, ai_count), 6),
        "authentic_recall": round(true_authentic / max(1, non_ai_count), 6),
        "inconclusive_rate": round(float(np.mean(inconclusive)), 6),
        "decisive_coverage": round(decisive / max(1, len(ai_truth)), 6),
        "decisive_accuracy": round(decisive_correct / max(1, decisive), 6),
        "likely_ai_precision": round(true_ai / max(1, true_ai + false_ai), 6),
        "likely_authentic_precision": round(true_authentic / max(1, true_authentic + ai_mislabeled_authentic), 6),
    }


def _binary_auc(truth: np.ndarray, scores: np.ndarray) -> float:
    positives = scores[truth == 1]
    negatives = scores[truth == 0]
    if len(positives) == 0 or len(negatives) == 0:
        return 0.5
    ordered = np.argsort(scores, kind="mergesort")
    ranks = np.empty(len(scores), dtype=np.float64)
    ranks[ordered] = np.arange(1, len(scores) + 1, dtype=np.float64)
    unique_scores, inverse, counts = np.unique(scores, return_inverse=True, return_counts=True)
    del unique_scores
    for group_id, count in enumerate(counts):
        if count > 1:
            members = inverse == group_id
            ranks[members] = float(np.mean(ranks[members]))
    positive_rank_sum = float(np.sum(ranks[truth == 1]))
    return (positive_rank_sum - len(positives) * (len(positives) + 1) / 2.0) / (len(positives) * len(negatives))


def _expected_calibration_error(truth: np.ndarray, scores: np.ndarray, bins: int = 10) -> float:
    error = 0.0
    for lower in np.linspace(0.0, 1.0, bins, endpoint=False):
        upper = lower + 1.0 / bins
        members = (scores >= lower) & (scores < upper if upper < 1.0 else scores <= upper)
        if not np.any(members):
            continue
        error += float(np.mean(members)) * abs(float(np.mean(scores[members])) - float(np.mean(truth[members])))
    return error


if __name__ == "__main__":
    main()