""" ItsNotAI v2 - AI Image Detector Gradio app for Hugging Face Spaces Dual-head model: Multi-class + Binary classification """ import gradio as gr import torch import torch.nn as nn import json from PIL import Image from transformers import AutoModelForImageClassification, AutoImageProcessor from huggingface_hub import hf_hub_download # Model configuration MODEL_ID = "boluobobo/ItsNotAI-ai-detector-v2" # Load model and processor print("Loading model...") model = AutoModelForImageClassification.from_pretrained(MODEL_ID) processor = AutoImageProcessor.from_pretrained(MODEL_ID) model.eval() # Load source metadata try: meta_path = hf_hub_download(repo_id=MODEL_ID, filename="source_meta.json") with open(meta_path) as f: meta = json.load(f) source_names = meta["source_names"] source_is_real = meta["source_is_real"] dual_head = meta.get("dual_head", False) hidden_size = meta.get("hidden_size", model.config.hidden_size) except Exception: # Fallback source_names = list(model.config.id2label.values()) source_is_real = {} dual_head = False hidden_size = model.config.hidden_size # Load binary head if dual-head mode binary_head = None if dual_head: try: binary_head_path = hf_hub_download(repo_id=MODEL_ID, filename="binary_head.pt") binary_head = nn.Sequential( nn.Dropout(0.1), nn.Linear(hidden_size, 2), ) binary_head.load_state_dict(torch.load(binary_head_path, map_location="cpu", weights_only=True)) binary_head.eval() print("Binary head loaded - using dual-head mode") except Exception as e: print(f"Binary head not found, falling back to single-head mode: {e}") dual_head = False print(f"Loaded {len(source_names)} classes, dual_head={dual_head}") def get_backbone_features(pixel_values): """Extract backbone features (CLS token)""" if hasattr(model, 'beit'): outputs = model.beit(pixel_values) return outputs.last_hidden_state[:, 0] elif hasattr(model, 'vit'): outputs = model.vit(pixel_values) return outputs.last_hidden_state[:, 0] else: for name, module in model.named_children(): if name not in ["classifier", "head", "fc"]: outputs = module(pixel_values) if hasattr(outputs, 'pooler_output') and outputs.pooler_output is not None: return outputs.pooler_output else: return outputs.last_hidden_state[:, 0] raise ValueError("Cannot extract backbone features from this model") def predict(image: Image.Image): """Predict if image is real or AI-generated""" if image is None: return None, None, "Please upload an image", None # Preprocess image = image.convert("RGB") inputs = processor(image, return_tensors="pt") # Inference with torch.no_grad(): outputs = model(**inputs) probs = torch.softmax(outputs.logits, dim=-1)[0] # Dual-head mode: use binary head if dual_head and binary_head is not None: features = get_backbone_features(inputs["pixel_values"]) binary_logits = binary_head(features) binary_probs = torch.softmax(binary_logits, dim=-1)[0] human_prob = binary_probs[0].item() # index 0 = Real ai_prob = binary_probs[1].item() # index 1 = AI is_real = human_prob > ai_prob else: # Fallback: Top-1 decision pred_idx = probs.argmax().item() predicted_source_top1 = source_names[pred_idx] confidence = probs[pred_idx].item() is_real = source_is_real.get(predicted_source_top1, False) if is_real: human_prob = confidence ai_prob = 1.0 - human_prob else: ai_prob = confidence human_prob = 1.0 - ai_prob # Top-1 multi-class prediction (for source display) pred_idx = probs.argmax().item() predicted_source = source_names[pred_idx] confidence = probs[pred_idx].item() # Get top 3 AI sources only (exclude real sources) ai_sources = [] for i, (name, prob) in enumerate(zip(source_names, probs.tolist())): if not source_is_real.get(name, False): ai_sources.append({"label": name, "score": round(prob, 3)}) ai_sources.sort(key=lambda x: x["score"], reverse=True) top3_sources = ai_sources[:3] # API-style JSON output api_output = { "ai_probability": round(ai_prob, 3), "human_probability": round(human_prob, 3), "predicted_source": predicted_source, "top3_sources": top3_sources, "dual_head": dual_head, } # Top predictions for bar chart top_preds = {} for i, (name, prob) in enumerate(zip(source_names, probs.tolist())): if prob > 0.01: marker = "[Real]" if source_is_real.get(name, False) else "[AI]" top_preds[f"{marker} {name}"] = prob top_preds = dict(sorted(top_preds.items(), key=lambda x: x[1], reverse=True)[:10]) # Summary mode_note = " *(Binary Head)*" if dual_head else "" verdict_emoji = "Human" if is_real else "AI Generated" verdict_color = "green" if is_real else "red" summary = f""" ### Verdict: **{verdict_emoji}**{mode_note} | Category | Probability | |----------|-------------| | Human | {human_prob:.1%} | | AI Generated | {ai_prob:.1%} | --- **Predicted Source**: {predicted_source} ({confidence:.1%}) **Top AI Sources**: """ for src in top3_sources[:3]: summary += f"\n- {src['label']}: {src['score']:.1%}" return ( {"Human": human_prob, "AI Generated": ai_prob}, top_preds, summary, api_output ) # Custom CSS css = """ .main-title { text-align: center; margin-bottom: 1rem; } .gradio-container { max-width: 1200px !important; } """ # Gradio interface with gr.Blocks(css=css, title="ItsNotAI v2 - AI Image Detector") as demo: gr.Markdown( """ # ItsNotAI v2 - AI Image Detector Upload an image to detect if it's **Human-made** or **AI-generated**, and identify the potential source. **New in v2**: Dual-head architecture with 95%+ binary accuracy. Supports FLUX, Midjourney, Stable Diffusion, DALL-E, and more. """, elem_classes="main-title" ) with gr.Row(): with gr.Column(scale=1): input_image = gr.Image( type="pil", label="Upload Image", height=400 ) submit_btn = gr.Button("Analyze", variant="primary", size="lg") with gr.Column(scale=1): # Main result result_label = gr.Label( label="Human vs AI", num_top_classes=2 ) # Top predictions top_preds_label = gr.Label( label="Top Predictions by Source", num_top_classes=10 ) # Detailed summary summary_md = gr.Markdown(label="Details") # API-style JSON output with gr.Accordion("API Output (JSON)", open=False): json_output = gr.JSON(label="API Output") # Event handlers submit_btn.click( fn=predict, inputs=[input_image], outputs=[result_label, top_preds_label, summary_md, json_output] ) input_image.change( fn=predict, inputs=[input_image], outputs=[result_label, top_preds_label, summary_md, json_output] ) gr.Markdown( """ --- ### About ItsNotAI Most AI detectors focus on catching AI. **We help artists prove their work is human-made.** - **Model**: [boluobobo/ItsNotAI-ai-detector-v2](https://huggingface.co/boluobobo/ItsNotAI-ai-detector-v2) - **Binary Accuracy**: 95.07% - **Multi-class Accuracy**: 93.47% - **Website**: [itsnotai.org](https://itsnotai.org) ### Disclaimer This tool is for educational and research purposes. Results should not be used as definitive proof of image authenticity. """ ) if __name__ == "__main__": demo.launch()