| from __future__ import annotations |
|
|
| import csv |
| import json |
| from pathlib import Path |
| from typing import Any, Optional |
|
|
| import typer |
| from PIL import Image, ImageOps |
| from rich.console import Console |
| from rich.table import Table |
| from tqdm import tqdm |
|
|
| from .config import ( |
| DEFAULT_BACKEND, |
| DEFAULT_HF_MODEL_ID, |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| DEFAULT_MODEL_NAME, |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| DEFAULT_PRETRAINED, |
| ) |
| from .evaluation import ( |
| build_combined_predictions, |
| search_blend_weight_threshold, |
| collect_folder_samples, |
| collect_predictions, |
| collect_tiny_genimage_parquet_samples, |
| combine_scores, |
| compute_metrics, |
| evaluate_folder, |
| evaluate_hf_dataset, |
| group_prediction_rows_against_reference, |
| group_prediction_rows, |
| search_hybrid_weight_threshold, |
| search_threshold, |
| split_samples_balanced, |
| write_report, |
| ) |
| from .model import create_detector, iter_images |
| from .types import DetectionResult, EvaluationReport |
|
|
| app = typer.Typer(help="Simple AI-generated image detector powered by UnivFD/CLIP.") |
| console = Console() |
| VALID_OPTIMIZE_METRICS = {"balanced_accuracy", "f1", "recall", "precision"} |
|
|
|
|
| @app.command() |
| def detect( |
| path: Path = typer.Argument(..., help="Image file or directory."), |
| recursive: bool = typer.Option(True, "--recursive/--no-recursive", help="Scan folders recursively."), |
| threshold: float = typer.Option(0.5, help="AI probability threshold."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| DEFAULT_BACKEND, |
| help="Detector backend: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| model_name: str = typer.Option(DEFAULT_MODEL_NAME, help="OpenCLIP model name for UnivFD."), |
| pretrained: str = typer.Option(DEFAULT_PRETRAINED, help="OpenCLIP pretrained tag for UnivFD."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for --backend hf."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for UnivFD when --backend hybrid.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid when --backend hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid-plus when --backend ultra.", |
| ), |
| batch_size: int = typer.Option(16, min=1, help="Batch size for inference."), |
| json_output: bool = typer.Option(False, "--json", help="Print JSON lines instead of a table."), |
| csv_output: Optional[Path] = typer.Option(None, "--csv", help="Save CSV report."), |
| ) -> None: |
| """Detect whether an image, or images in a folder, are AI-generated.""" |
| paths = iter_images(path, recursive=recursive) |
| if not paths: |
| raise typer.BadParameter(f"No supported images found at {path}") |
|
|
| detector = create_detector( |
| backend, |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| results = _predict_paths(detector, paths, batch_size=batch_size, quiet=json_output) |
|
|
| if csv_output: |
| _write_csv(results, csv_output) |
| console.print(f"Saved CSV report to {csv_output}") |
|
|
| if json_output: |
| for result in results: |
| print(json.dumps(result.as_dict(), ensure_ascii=False)) |
| return |
|
|
| _print_detection_table(results) |
| console.print("[yellow]Note:[/yellow] AI image detection is probabilistic; do not use it as sole proof.") |
|
|
|
|
| @app.command("benchmark-folder") |
| def benchmark_folder_command( |
| root: Path = typer.Argument(..., help="Dataset root containing real/fake folders."), |
| real_dir: str = typer.Option("real", help="Folder name containing real images."), |
| fake_dir: str = typer.Option("ai", help="Folder name containing AI/fake images."), |
| recursive: bool = typer.Option(True, "--recursive/--no-recursive", help="Scan folders recursively."), |
| max_per_class: Optional[int] = typer.Option(None, min=1, help="Optional cap per class."), |
| threshold: float = typer.Option(0.5, help="AI probability threshold."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| DEFAULT_BACKEND, |
| help="Detector backend: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| model_name: str = typer.Option(DEFAULT_MODEL_NAME, help="OpenCLIP model name for UnivFD."), |
| pretrained: str = typer.Option(DEFAULT_PRETRAINED, help="OpenCLIP pretrained tag for UnivFD."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for --backend hf."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for UnivFD when --backend hybrid.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid when --backend hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid-plus when --backend ultra.", |
| ), |
| batch_size: int = typer.Option(16, min=1, help="Batch size for inference."), |
| output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write JSON benchmark report."), |
| ) -> None: |
| """Evaluate on an image-folder benchmark.""" |
| detector = create_detector( |
| backend, |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| report = evaluate_folder( |
| detector, |
| root, |
| real_dir=real_dir, |
| fake_dir=fake_dir, |
| recursive=recursive, |
| max_per_class=max_per_class, |
| batch_size=batch_size, |
| ) |
| _finish_benchmark(report, output) |
|
|
|
|
| @app.command("benchmark-hf") |
| def benchmark_hf_command( |
| dataset_name: str = typer.Argument(..., help="Hugging Face dataset name."), |
| split: str = typer.Option("validation", help="Dataset split."), |
| image_field: str = typer.Option("image", help="Image column name."), |
| label_field: str = typer.Option("label", help="Binary label column name."), |
| fake_label: str = typer.Option("1", help="Value in label column that means AI/fake."), |
| max_samples: Optional[int] = typer.Option(None, min=1, help="Optional sample cap."), |
| streaming: bool = typer.Option(False, help="Use streaming dataset loading."), |
| shuffle_seed: Optional[int] = typer.Option(None, help="Shuffle before sampling."), |
| trust_remote_code: bool = typer.Option(False, help="Allow remote dataset code."), |
| threshold: float = typer.Option(0.5, help="AI probability threshold."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| DEFAULT_BACKEND, |
| help="Detector backend: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| model_name: str = typer.Option(DEFAULT_MODEL_NAME, help="OpenCLIP model name for UnivFD."), |
| pretrained: str = typer.Option(DEFAULT_PRETRAINED, help="OpenCLIP pretrained tag for UnivFD."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for --backend hf."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for UnivFD when --backend hybrid.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid when --backend hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid-plus when --backend ultra.", |
| ), |
| batch_size: int = typer.Option(16, min=1, help="Batch size for inference."), |
| output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write JSON benchmark report."), |
| ) -> None: |
| """Evaluate on a Hugging Face image dataset.""" |
| detector = create_detector( |
| backend, |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| report = evaluate_hf_dataset( |
| detector, |
| dataset_name, |
| split=split, |
| image_field=image_field, |
| label_field=label_field, |
| fake_label=fake_label, |
| max_samples=max_samples, |
| batch_size=batch_size, |
| streaming=streaming, |
| shuffle_seed=shuffle_seed, |
| trust_remote_code=trust_remote_code, |
| ) |
| _finish_benchmark(report, output) |
|
|
|
|
| @app.command("benchmark-calibrated-folder") |
| def benchmark_calibrated_folder_command( |
| root: Path = typer.Argument(..., help="Dataset root containing real/fake folders."), |
| real_dir: str = typer.Option("real", help="Folder name containing real images."), |
| fake_dir: str = typer.Option("ai", help="Folder name containing AI/fake images."), |
| recursive: bool = typer.Option(True, "--recursive/--no-recursive", help="Scan folders recursively."), |
| max_per_class: Optional[int] = typer.Option(None, min=2, help="Optional cap per class."), |
| calibration_fraction: float = typer.Option( |
| 0.5, |
| min=0.05, |
| max=0.95, |
| help="Fraction of each class reserved for calibration.", |
| ), |
| split_seed: str = typer.Option("aidetect", help="Seed used for deterministic calibration/test split."), |
| threshold: float = typer.Option(0.5, help="Fallback threshold before calibration."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| "hybrid", |
| help="Backend to calibrate: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| model_name: str = typer.Option(DEFAULT_MODEL_NAME, help="OpenCLIP model name for UnivFD."), |
| pretrained: str = typer.Option(DEFAULT_PRETRAINED, help="OpenCLIP pretrained tag for UnivFD."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for HF or hybrid."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Initial blend weight for UnivFD when backend=hybrid.", |
| ), |
| hybrid_alpha_step: float = typer.Option( |
| 0.05, |
| min=0.01, |
| max=0.5, |
| help="Grid step for hybrid weight search during calibration.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Initial blend weight for hybrid when backend=hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Initial blend weight for hybrid-plus when backend=ultra.", |
| ), |
| optimize_metric: str = typer.Option( |
| "balanced_accuracy", |
| help="Calibration objective: balanced_accuracy, f1, recall, or precision.", |
| ), |
| min_recall: Optional[float] = typer.Option( |
| None, |
| min=0.0, |
| max=1.0, |
| help="Optional minimum calibration recall constraint for threshold search.", |
| ), |
| batch_size: int = typer.Option(16, min=1, help="Batch size for inference."), |
| output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write JSON calibration benchmark report."), |
| ) -> None: |
| """Calibrate threshold on a held-out split, then evaluate on a disjoint test split.""" |
| _validate_optimize_metric(optimize_metric) |
| samples = collect_folder_samples( |
| root, |
| real_dir=real_dir, |
| fake_dir=fake_dir, |
| recursive=recursive, |
| max_per_class=max_per_class, |
| ) |
| if len(samples) < 4: |
| raise typer.BadParameter("Need at least two real and two AI images for calibrated benchmarking.") |
|
|
| calibration_samples, test_samples = split_samples_balanced( |
| samples, |
| calibration_fraction=calibration_fraction, |
| seed=split_seed, |
| ) |
| report = _run_calibrated_folder_benchmark( |
| calibration_samples=calibration_samples, |
| test_samples=test_samples, |
| backend=backend, |
| threshold=threshold, |
| device=device, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| hybrid_alpha_step=hybrid_alpha_step, |
| optimize_metric=optimize_metric, |
| min_recall=min_recall, |
| batch_size=batch_size, |
| dataset_info={ |
| "kind": "folder-calibrated", |
| "root": str(root), |
| "real_dir": real_dir, |
| "fake_dir": fake_dir, |
| "recursive": recursive, |
| "max_per_class": max_per_class, |
| "calibration_fraction": calibration_fraction, |
| "split_seed": split_seed, |
| "calibration_size": len(calibration_samples), |
| "test_size": len(test_samples), |
| }, |
| ) |
| if output: |
| output.parent.mkdir(parents=True, exist_ok=True) |
| output.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8") |
| console.print(f"Saved calibrated benchmark report to {output}") |
| _print_calibrated_benchmark(report) |
|
|
|
|
| @app.command("benchmark-tiny-genimage-local") |
| def benchmark_tiny_genimage_local_command( |
| parquet_paths: list[Path] = typer.Argument(..., help="One or more local Tiny-GenImage parquet shards."), |
| calibration_fraction: float = typer.Option( |
| 0.5, |
| min=0.05, |
| max=0.95, |
| help="Fraction of each class reserved for calibration.", |
| ), |
| split_seed: str = typer.Option("aidetect", help="Seed used for deterministic calibration/test split."), |
| max_per_class_per_shard: Optional[int] = typer.Option( |
| None, |
| min=1, |
| help="Optional cap per class inside each parquet shard.", |
| ), |
| generators: Optional[str] = typer.Option( |
| None, |
| help="Comma-separated generator filter, for example 'Midjourney,SD14'.", |
| ), |
| threshold: float = typer.Option(0.5, help="Fallback threshold before calibration."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| "univfd", |
| help="Backend to calibrate: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| model_name: str = typer.Option(DEFAULT_MODEL_NAME, help="OpenCLIP model name for UnivFD."), |
| pretrained: str = typer.Option(DEFAULT_PRETRAINED, help="OpenCLIP pretrained tag for UnivFD."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for HF or hybrid."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Initial blend weight for UnivFD when backend=hybrid.", |
| ), |
| hybrid_alpha_step: float = typer.Option( |
| 0.05, |
| min=0.01, |
| max=0.5, |
| help="Grid step for hybrid weight search during calibration.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Initial blend weight for hybrid when backend=hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Initial blend weight for hybrid-plus when backend=ultra.", |
| ), |
| optimize_metric: str = typer.Option( |
| "balanced_accuracy", |
| help="Calibration objective: balanced_accuracy, f1, recall, or precision.", |
| ), |
| min_recall: Optional[float] = typer.Option( |
| None, |
| min=0.0, |
| max=1.0, |
| help="Optional minimum calibration recall constraint for threshold search.", |
| ), |
| batch_size: int = typer.Option(16, min=1, help="Batch size for inference."), |
| output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write JSON report."), |
| ) -> None: |
| """Benchmark one or more local Tiny-GenImage parquet shards with per-generator summaries.""" |
| _validate_optimize_metric(optimize_metric) |
| generator_filter = None |
| if generators: |
| generator_filter = {value.strip() for value in generators.split(",") if value.strip()} |
|
|
| samples = collect_tiny_genimage_parquet_samples( |
| parquet_paths, |
| max_per_class_per_shard=max_per_class_per_shard, |
| generators=generator_filter, |
| ) |
| if len(samples) < 4: |
| raise typer.BadParameter("Need at least two real and two AI samples after filtering.") |
|
|
| calibration_samples, test_samples = split_samples_balanced( |
| samples, |
| calibration_fraction=calibration_fraction, |
| seed=split_seed, |
| ) |
| report = _run_calibrated_folder_benchmark( |
| calibration_samples=calibration_samples, |
| test_samples=test_samples, |
| backend=backend, |
| threshold=threshold, |
| device=device, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| hybrid_alpha_step=hybrid_alpha_step, |
| optimize_metric=optimize_metric, |
| min_recall=min_recall, |
| batch_size=batch_size, |
| dataset_info={ |
| "kind": "tiny-genimage-local", |
| "parquet_paths": [str(path) for path in parquet_paths], |
| "calibration_fraction": calibration_fraction, |
| "split_seed": split_seed, |
| "max_per_class_per_shard": max_per_class_per_shard, |
| "generator_filter": sorted(generator_filter) if generator_filter else None, |
| "calibration_size": len(calibration_samples), |
| "test_size": len(test_samples), |
| }, |
| group_fields=("generator", "source_shard"), |
| ) |
| if output: |
| output.parent.mkdir(parents=True, exist_ok=True) |
| output.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8") |
| console.print(f"Saved Tiny-GenImage benchmark report to {output}") |
| _print_calibrated_benchmark(report) |
|
|
|
|
| @app.command("prepare-tiny-genimage") |
| def prepare_tiny_genimage_command( |
| output_dir: Path = typer.Argument(..., help="Output folder with real/ai subfolders."), |
| repo_id: str = typer.Option("TheKernel01/Tiny-GenImage", help="Hugging Face dataset repo."), |
| filename: str = typer.Option( |
| "data/validation-00000-of-00004.parquet", |
| help="Parquet shard inside the dataset repo.", |
| ), |
| local_parquet: Optional[Path] = typer.Option( |
| None, |
| help="Use an already-downloaded parquet shard instead of fetching from Hugging Face.", |
| ), |
| max_per_class: int = typer.Option(20, min=1, help="Images to export per class."), |
| overwrite: bool = typer.Option(False, help="Overwrite existing exported images."), |
| ) -> None: |
| """Download a Tiny-GenImage shard and export a balanced folder benchmark.""" |
| try: |
| from datasets import load_dataset |
| except ImportError as exc: |
| raise typer.BadParameter("Install eval extras first: pip install 'ai-image-detector[eval]'") from exc |
|
|
| if local_parquet is not None: |
| parquet_path = str(local_parquet) |
| else: |
| try: |
| from huggingface_hub import hf_hub_download |
| except ImportError as exc: |
| raise typer.BadParameter("Install huggingface_hub or pass --local-parquet.") from exc |
| parquet_path = hf_hub_download(repo_id=repo_id, repo_type="dataset", filename=filename) |
| dataset = load_dataset("parquet", data_files={"data": parquet_path}, split="data") |
| real_dir = output_dir / "real" |
| fake_dir = output_dir / "ai" |
| real_dir.mkdir(parents=True, exist_ok=True) |
| fake_dir.mkdir(parents=True, exist_ok=True) |
|
|
| exported = {"real": 0, "ai": 0} |
| label_names = getattr(dataset.features.get("label"), "names", None) |
| generator_names = getattr(dataset.features.get("generator"), "names", None) |
| for index, item in enumerate(dataset): |
| label_value = int(item["label"]) |
| class_name = _tiny_genimage_class_name(label_value, label_names) |
| if class_name not in exported or exported[class_name] >= max_per_class: |
| continue |
| generator = _safe_name(_class_label_name(item.get("generator"), generator_names)) |
| filename_out = f"{index:06d}-{generator}.jpg" |
| output_path = (real_dir if class_name == "real" else fake_dir) / filename_out |
| if output_path.exists() and not overwrite: |
| exported[class_name] += 1 |
| else: |
| item["image"].convert("RGB").save(output_path, quality=95) |
| exported[class_name] += 1 |
| if all(count >= max_per_class for count in exported.values()): |
| break |
|
|
| if any(count < max_per_class for count in exported.values()): |
| raise typer.BadParameter( |
| f"Could only export {exported}; try a different shard or lower --max-per-class." |
| ) |
| console.print( |
| f"Exported Tiny-GenImage folder benchmark to {output_dir} " |
| f"({exported['real']} real, {exported['ai']} ai)." |
| ) |
|
|
|
|
| @app.command() |
| def serve( |
| threshold: float = typer.Option(0.5, help="AI probability threshold."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| DEFAULT_BACKEND, |
| help="Detector backend: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for --backend hf."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for UnivFD when --backend hybrid.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid when --backend hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid-plus when --backend ultra.", |
| ), |
| ) -> None: |
| """Launch a local Gradio web UI.""" |
| try: |
| import gradio as gr |
| except ImportError as exc: |
| raise typer.BadParameter("Install web extras first: pip install 'ai-image-detector[web]'") from exc |
|
|
| detector = create_detector( |
| backend, |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
|
|
| def predict(image): |
| if image is None: |
| return {"error": "Please upload an image."} |
| result = detector.predict_image(image) |
| return result.as_dict() |
|
|
| demo = gr.Interface( |
| fn=predict, |
| inputs=gr.Image(type="pil", label="Upload image"), |
| outputs=gr.JSON(label="Detection result"), |
| title="AI Image Detector", |
| description="AI-generated image detector. Results are probabilistic.", |
| ) |
| demo.launch() |
|
|
|
|
| @app.command("api") |
| def api_command( |
| host: str = typer.Option("127.0.0.1", help="Host to bind."), |
| port: int = typer.Option(8000, help="Port to bind."), |
| threshold: float = typer.Option(0.5, help="AI probability threshold."), |
| device: Optional[str] = typer.Option(None, help="cuda, mps, cpu, or auto when omitted."), |
| backend: str = typer.Option( |
| DEFAULT_BACKEND, |
| help="Detector backend: univfd, hf, nonescape-mini, nonescape-full, sentry-convnext-small, hybrid, hybrid-plus, or ultra.", |
| ), |
| weight_path: Optional[Path] = typer.Option(None, help="Optional local UnivFD fc_weights.pth path."), |
| hf_model: str = typer.Option(DEFAULT_HF_MODEL_ID, help="Hugging Face model id for --backend hf."), |
| hybrid_univfd_weight: float = typer.Option( |
| DEFAULT_HYBRID_UNIVFD_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for UnivFD when --backend hybrid.", |
| ), |
| hybrid_plus_primary_weight: float = typer.Option( |
| DEFAULT_HYBRID_PLUS_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid when --backend hybrid-plus.", |
| ), |
| ultra_primary_weight: float = typer.Option( |
| DEFAULT_ULTRA_PRIMARY_WEIGHT, |
| min=0.0, |
| max=1.0, |
| help="Blend weight for hybrid-plus when --backend ultra.", |
| ), |
| ) -> None: |
| """Launch a FastAPI service.""" |
| try: |
| import uvicorn |
| except ImportError as exc: |
| raise typer.BadParameter("Install API extras first: pip install 'ai-image-detector[api]'") from exc |
|
|
| from .api import create_app |
|
|
| fastapi_app = create_app( |
| threshold=threshold, |
| device=device, |
| backend=backend, |
| weight_path=weight_path, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| uvicorn.run(fastapi_app, host=host, port=port) |
|
|
|
|
| def _predict_paths(detector, paths: list[Path], *, batch_size: int, quiet: bool) -> list[DetectionResult]: |
| results: list[DetectionResult] = [] |
| batch_images = [] |
| batch_paths = [] |
| iterator = tqdm(paths, disable=quiet or len(paths) == 1) |
| for image_path in iterator: |
| try: |
| with Image.open(image_path) as image: |
| batch_images.append(ImageOps.exif_transpose(image).convert("RGB")) |
| batch_paths.append(image_path) |
| if len(batch_images) >= batch_size: |
| results.extend(detector.predict_images(batch_images, paths=batch_paths)) |
| batch_images = [] |
| batch_paths = [] |
| except Exception as exc: |
| console.print(f"[red]Failed:[/red] {image_path} ({exc})") |
| if batch_images: |
| results.extend(detector.predict_images(batch_images, paths=batch_paths)) |
| return results |
|
|
|
|
| def _write_csv(results: list[DetectionResult], output: Path) -> None: |
| output.parent.mkdir(parents=True, exist_ok=True) |
| with output.open("w", newline="", encoding="utf-8") as f: |
| fieldnames = [ |
| "path", |
| "label", |
| "probability_ai", |
| "probability_real", |
| "confidence", |
| "raw_score", |
| "backend", |
| ] |
| writer = csv.DictWriter(f, fieldnames=fieldnames) |
| writer.writeheader() |
| writer.writerows([r.as_dict() for r in results]) |
|
|
|
|
| def _print_detection_table(results: list[DetectionResult]) -> None: |
| table = Table(title="AI Image Detection Results") |
| table.add_column("Image") |
| table.add_column("Label") |
| table.add_column("AI Prob", justify="right") |
| table.add_column("Real Prob", justify="right") |
| table.add_column("Confidence", justify="right") |
| table.add_column("Backend") |
| for result in results: |
| row_style = "red" if result.label == "ai" else "green" |
| table.add_row( |
| str(result.path), |
| result.label.upper(), |
| f"{result.probability_ai:.3f}", |
| f"{result.probability_real:.3f}", |
| f"{result.confidence:.3f}", |
| result.backend, |
| style=row_style, |
| ) |
| console.print(table) |
|
|
|
|
| def _finish_benchmark(report: EvaluationReport, output: Path | None) -> None: |
| if output: |
| write_report(report, output) |
| console.print(f"Saved benchmark report to {output}") |
| metrics = report.metrics |
| table = Table(title="Benchmark Results") |
| table.add_column("Dataset") |
| table.add_column("N", justify="right") |
| table.add_column("Accuracy", justify="right") |
| table.add_column("Balanced Acc", justify="right") |
| table.add_column("F1", justify="right") |
| table.add_column("ROC AUC", justify="right") |
| table.add_column("Images/s", justify="right") |
| table.add_row( |
| metrics.dataset, |
| str(metrics.n_samples), |
| f"{metrics.accuracy:.3f}", |
| f"{metrics.balanced_accuracy:.3f}", |
| f"{metrics.f1:.3f}", |
| "n/a" if metrics.roc_auc is None else f"{metrics.roc_auc:.3f}", |
| f"{metrics.images_per_second:.2f}", |
| ) |
| console.print(table) |
|
|
|
|
| def _run_calibrated_folder_benchmark( |
| *, |
| calibration_samples, |
| test_samples, |
| backend: str, |
| threshold: float, |
| device: Optional[str], |
| weight_path: Optional[Path], |
| model_name: str, |
| pretrained: str, |
| hf_model: str, |
| hybrid_univfd_weight: float, |
| hybrid_plus_primary_weight: float, |
| ultra_primary_weight: float, |
| hybrid_alpha_step: float, |
| optimize_metric: str, |
| min_recall: Optional[float], |
| batch_size: int, |
| dataset_info: dict[str, Any], |
| group_fields: tuple[str, ...] = (), |
| ) -> dict[str, Any]: |
| backend = backend.lower() |
| if backend in {"ultra", "hybrid-ultra", "sentry-plus"}: |
| primary_detector = create_detector( |
| "hybrid-plus", |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| secondary_detector = create_detector( |
| "sentry-convnext-small", |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| calibration_primary = collect_predictions(primary_detector, calibration_samples, batch_size=batch_size) |
| calibration_secondary = collect_predictions(secondary_detector, calibration_samples, batch_size=batch_size) |
| search = search_blend_weight_threshold( |
| calibration_primary.y_true, |
| calibration_primary.y_score, |
| calibration_secondary.y_score, |
| alpha_step=hybrid_alpha_step, |
| objective=optimize_metric, |
| min_recall=min_recall, |
| ) |
| selected_weight = float(search["primary_weight"] or ultra_primary_weight) |
| selected_threshold = float(search["threshold"] or threshold) |
| calibration_scores = combine_scores( |
| calibration_primary.y_score, |
| calibration_secondary.y_score, |
| univfd_weight=selected_weight, |
| ) |
| test_primary = collect_predictions(primary_detector, test_samples, batch_size=batch_size) |
| test_secondary = collect_predictions(secondary_detector, test_samples, batch_size=batch_size) |
| test_scores = combine_scores( |
| test_primary.y_score, |
| test_secondary.y_score, |
| univfd_weight=selected_weight, |
| ) |
| calibration_metrics = compute_metrics( |
| calibration_primary.y_true, |
| calibration_scores, |
| dataset="calibration", |
| threshold=selected_threshold, |
| seconds=calibration_primary.seconds + calibration_secondary.seconds, |
| ) |
| test_metrics = compute_metrics( |
| test_primary.y_true, |
| test_scores, |
| dataset="test", |
| threshold=selected_threshold, |
| seconds=test_primary.seconds + test_secondary.seconds, |
| ) |
| calibration_predictions = build_combined_predictions( |
| calibration_primary.predictions, |
| calibration_scores, |
| threshold=selected_threshold, |
| backend="ultra", |
| ) |
| test_predictions = build_combined_predictions( |
| test_primary.predictions, |
| test_scores, |
| threshold=selected_threshold, |
| backend="ultra", |
| ) |
| model = { |
| "backend": "ultra", |
| "selected_threshold": selected_threshold, |
| "selected_primary_weight": selected_weight, |
| "search_alpha_step": hybrid_alpha_step, |
| "optimize_metric": optimize_metric, |
| "min_recall": min_recall, |
| "components": { |
| "primary": primary_detector.model_info(), |
| "secondary": secondary_detector.model_info(), |
| }, |
| } |
| search_summary = search |
| elif backend in {"hybrid-plus", "hybrid_plus", "stacked"}: |
| primary_detector = create_detector( |
| "hybrid", |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| secondary_detector = create_detector( |
| "nonescape-mini", |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| calibration_primary = collect_predictions(primary_detector, calibration_samples, batch_size=batch_size) |
| calibration_secondary = collect_predictions(secondary_detector, calibration_samples, batch_size=batch_size) |
| search = search_blend_weight_threshold( |
| calibration_primary.y_true, |
| calibration_primary.y_score, |
| calibration_secondary.y_score, |
| alpha_step=hybrid_alpha_step, |
| objective=optimize_metric, |
| min_recall=min_recall, |
| ) |
| selected_weight = float(search["primary_weight"] or hybrid_plus_primary_weight) |
| selected_threshold = float(search["threshold"] or threshold) |
| calibration_scores = combine_scores( |
| calibration_primary.y_score, |
| calibration_secondary.y_score, |
| univfd_weight=selected_weight, |
| ) |
| test_primary = collect_predictions(primary_detector, test_samples, batch_size=batch_size) |
| test_secondary = collect_predictions(secondary_detector, test_samples, batch_size=batch_size) |
| test_scores = combine_scores( |
| test_primary.y_score, |
| test_secondary.y_score, |
| univfd_weight=selected_weight, |
| ) |
| calibration_metrics = compute_metrics( |
| calibration_primary.y_true, |
| calibration_scores, |
| dataset="calibration", |
| threshold=selected_threshold, |
| seconds=calibration_primary.seconds + calibration_secondary.seconds, |
| ) |
| test_metrics = compute_metrics( |
| test_primary.y_true, |
| test_scores, |
| dataset="test", |
| threshold=selected_threshold, |
| seconds=test_primary.seconds + test_secondary.seconds, |
| ) |
| calibration_predictions = build_combined_predictions( |
| calibration_primary.predictions, |
| calibration_scores, |
| threshold=selected_threshold, |
| backend="hybrid-plus", |
| ) |
| test_predictions = build_combined_predictions( |
| test_primary.predictions, |
| test_scores, |
| threshold=selected_threshold, |
| backend="hybrid-plus", |
| ) |
| model = { |
| "backend": "hybrid-plus", |
| "selected_threshold": selected_threshold, |
| "selected_primary_weight": selected_weight, |
| "search_alpha_step": hybrid_alpha_step, |
| "optimize_metric": optimize_metric, |
| "min_recall": min_recall, |
| "components": { |
| "primary": primary_detector.model_info(), |
| "secondary": secondary_detector.model_info(), |
| }, |
| } |
| search_summary = search |
| elif backend in {"hybrid", "ensemble"}: |
| univfd_detector = create_detector( |
| "univfd", |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| hf_detector = create_detector( |
| "hf", |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| calibration_univfd = collect_predictions(univfd_detector, calibration_samples, batch_size=batch_size) |
| calibration_hf = collect_predictions(hf_detector, calibration_samples, batch_size=batch_size) |
| search = search_hybrid_weight_threshold( |
| calibration_univfd.y_true, |
| calibration_univfd.y_score, |
| calibration_hf.y_score, |
| alpha_step=hybrid_alpha_step, |
| objective=optimize_metric, |
| min_recall=min_recall, |
| ) |
| selected_weight = float(search["univfd_weight"] or hybrid_univfd_weight) |
| selected_threshold = float(search["threshold"] or threshold) |
| calibration_scores = combine_scores( |
| calibration_univfd.y_score, |
| calibration_hf.y_score, |
| univfd_weight=selected_weight, |
| ) |
| test_univfd = collect_predictions(univfd_detector, test_samples, batch_size=batch_size) |
| test_hf = collect_predictions(hf_detector, test_samples, batch_size=batch_size) |
| test_scores = combine_scores( |
| test_univfd.y_score, |
| test_hf.y_score, |
| univfd_weight=selected_weight, |
| ) |
| calibration_metrics = compute_metrics( |
| calibration_univfd.y_true, |
| calibration_scores, |
| dataset="calibration", |
| threshold=selected_threshold, |
| seconds=calibration_univfd.seconds + calibration_hf.seconds, |
| ) |
| test_metrics = compute_metrics( |
| test_univfd.y_true, |
| test_scores, |
| dataset="test", |
| threshold=selected_threshold, |
| seconds=test_univfd.seconds + test_hf.seconds, |
| ) |
| calibration_predictions = build_combined_predictions( |
| calibration_univfd.predictions, |
| calibration_scores, |
| threshold=selected_threshold, |
| backend="hybrid", |
| ) |
| test_predictions = build_combined_predictions( |
| test_univfd.predictions, |
| test_scores, |
| threshold=selected_threshold, |
| backend="hybrid", |
| ) |
| model = { |
| "backend": "hybrid", |
| "selected_threshold": selected_threshold, |
| "selected_univfd_weight": selected_weight, |
| "search_alpha_step": hybrid_alpha_step, |
| "optimize_metric": optimize_metric, |
| "min_recall": min_recall, |
| "components": { |
| "univfd": univfd_detector.model_info(), |
| "hf": hf_detector.model_info(), |
| }, |
| } |
| search_summary = search |
| else: |
| detector = create_detector( |
| backend, |
| device=device, |
| threshold=threshold, |
| weight_path=weight_path, |
| model_name=model_name, |
| pretrained=pretrained, |
| hf_model=hf_model, |
| hybrid_univfd_weight=hybrid_univfd_weight, |
| hybrid_plus_primary_weight=hybrid_plus_primary_weight, |
| ultra_primary_weight=ultra_primary_weight, |
| ) |
| calibration_run = collect_predictions(detector, calibration_samples, batch_size=batch_size) |
| search = search_threshold( |
| calibration_run.y_true, |
| calibration_run.y_score, |
| objective=optimize_metric, |
| min_recall=min_recall, |
| ) |
| selected_threshold = float(search["threshold"] or threshold) |
| test_run = collect_predictions(detector, test_samples, batch_size=batch_size) |
| calibration_metrics = compute_metrics( |
| calibration_run.y_true, |
| calibration_run.y_score, |
| dataset="calibration", |
| threshold=selected_threshold, |
| seconds=calibration_run.seconds, |
| ) |
| test_metrics = compute_metrics( |
| test_run.y_true, |
| test_run.y_score, |
| dataset="test", |
| threshold=selected_threshold, |
| seconds=test_run.seconds, |
| ) |
| calibration_predictions = build_combined_predictions( |
| calibration_run.predictions, |
| calibration_run.y_score, |
| threshold=selected_threshold, |
| backend=backend, |
| ) |
| test_predictions = build_combined_predictions( |
| test_run.predictions, |
| test_run.y_score, |
| threshold=selected_threshold, |
| backend=backend, |
| ) |
| model = detector.model_info() | { |
| "selected_threshold": selected_threshold, |
| "optimize_metric": optimize_metric, |
| "min_recall": min_recall, |
| } |
| search_summary = search |
|
|
| report = { |
| "mode": "calibrated-folder", |
| "dataset": dataset_info, |
| "model": model, |
| "calibration": { |
| "metrics": calibration_metrics.as_dict(), |
| "predictions": calibration_predictions, |
| }, |
| "test": { |
| "metrics": test_metrics.as_dict(), |
| "predictions": test_predictions, |
| }, |
| "search": search_summary, |
| } |
| if group_fields: |
| threshold_value = float(report["test"]["metrics"]["threshold"]) |
| report["groups"] = {} |
| for field in group_fields: |
| report["groups"][field] = group_prediction_rows( |
| report["test"]["predictions"], |
| field=field, |
| threshold=threshold_value, |
| ) |
| if field == "generator": |
| report["groups"]["generator_vs_real"] = group_prediction_rows_against_reference( |
| report["test"]["predictions"], |
| field="generator", |
| reference_value="Real", |
| threshold=threshold_value, |
| ) |
| return report |
|
|
|
|
| def _print_calibrated_benchmark(report: dict[str, Any]) -> None: |
| table = Table(title="Calibrated Benchmark Results") |
| table.add_column("Split") |
| table.add_column("N", justify="right") |
| table.add_column("Threshold", justify="right") |
| table.add_column("Accuracy", justify="right") |
| table.add_column("Balanced Acc", justify="right") |
| table.add_column("F1", justify="right") |
| table.add_column("ROC AUC", justify="right") |
| table.add_column("Images/s", justify="right") |
| for split in ("calibration", "test"): |
| metrics = report[split]["metrics"] |
| table.add_row( |
| split, |
| str(metrics["n_samples"]), |
| f"{metrics['threshold']:.3f}", |
| f"{metrics['accuracy']:.3f}", |
| f"{metrics['balanced_accuracy']:.3f}", |
| f"{metrics['f1']:.3f}", |
| "n/a" if metrics["roc_auc"] is None else f"{metrics['roc_auc']:.3f}", |
| f"{metrics['images_per_second']:.2f}", |
| ) |
| console.print(table) |
| model = report["model"] |
| if model.get("backend") == "hybrid": |
| console.print( |
| "Selected hybrid setting: " |
| f"threshold={model['selected_threshold']:.3f}, " |
| f"univfd_weight={model['selected_univfd_weight']:.3f}" |
| ) |
| if model.get("backend") == "hybrid-plus": |
| console.print( |
| "Selected hybrid-plus setting: " |
| f"threshold={model['selected_threshold']:.3f}, " |
| f"hybrid_weight={model['selected_primary_weight']:.3f}" |
| ) |
| if model.get("backend") == "ultra": |
| console.print( |
| "Selected ultra setting: " |
| f"threshold={model['selected_threshold']:.3f}, " |
| f"hybrid-plus_weight={model['selected_primary_weight']:.3f}" |
| ) |
|
|
|
|
| def _tiny_genimage_class_name(label_value: int, label_names: list[str] | None) -> str: |
| if label_names and 0 <= label_value < len(label_names): |
| label_name = label_names[label_value].lower() |
| if label_name in {"fake", "ai", "synthetic", "generated"}: |
| return "ai" |
| if label_name in {"real", "natural"}: |
| return "real" |
| return "ai" if label_value == 1 else "real" |
|
|
|
|
| def _class_label_name(value, names: list[str] | None) -> str: |
| if names is not None and value is not None: |
| index = int(value) |
| if 0 <= index < len(names): |
| return names[index] |
| return "unknown" if value is None else str(value) |
|
|
|
|
| def _safe_name(value: str) -> str: |
| return "".join(character if character.isalnum() or character in "-_" else "-" for character in value) |
|
|
|
|
| def _validate_optimize_metric(value: str) -> None: |
| if value not in VALID_OPTIMIZE_METRICS: |
| raise typer.BadParameter( |
| f"Unsupported optimize metric: {value!r}. Choose from {sorted(VALID_OPTIMIZE_METRICS)}." |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| app() |
|
|
