Spaces:
Sleeping
Sleeping
File size: 2,160 Bytes
825f4e1 780829b 825f4e1 780829b 825f4e1 780829b 825f4e1 780829b 825f4e1 780829b 825f4e1 20f32b6 780829b 825f4e1 780829b 825f4e1 20f32b6 825f4e1 20f32b6 780829b 825f4e1 20f32b6 825f4e1 780829b 825f4e1 780829b 825f4e1 20f32b6 780829b 825f4e1 780829b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import gradio as gr
from transformers import pipeline
from PIL import Image
import traceback
# The model we will use
MODEL_ID = "umm-maybe/AI-image-detector"
# Load the model pipeline (image-classification)
pipe = pipeline("image-classification", model=MODEL_ID)
def predict_image(image: Image.Image):
"""
image: PIL image from Gradio
We focus on the fine-tuned ai_detector output if available.
"""
try:
# run the pipeline
res = pipe(image)
# typical pipeline returns list of dicts: [{"label": "...", "score": ...}, ...]
if isinstance(res, list) and isinstance(res[0], dict) and "label" in res[0]:
top = res[0]
label = top.get("label", "")
score = float(top.get("score", 0.0))
if label.lower().startswith("artificial") or "ai" in label.lower():
verdict = "AI-generated"
elif label.lower().startswith("human") or "real" in label.lower():
verdict = "Human-made"
else:
verdict = label
return f"{verdict} — {score:.2f}"
# Handle nested response structures
if isinstance(res, dict):
imr = res.get("individual_model_results") or res.get("models") or {}
ai_det = imr.get("ai_detector") or None
if ai_det:
ai_prob = ai_det.get("ai_probability") or ai_det.get("score") or ai_det.get("confidence")
ai_prob = float(ai_prob)
if ai_prob >= 0.6:
return f"AI-generated — {ai_prob:.2f}"
elif ai_prob <= 0.4:
return f"Human-made — {(1 - ai_prob):.2f}"
else:
return f"Uncertain — {ai_prob:.2f}"
return str(res)
except Exception as e:
traceback.print_exc()
return f"Error: {e}"
# Build Gradio app
demo = gr.Interface(
fn=predict_image,
inputs=gr.Image(type="pil"),
outputs="text",
title="AI Image Detector",
description="Upload an image to detect whether it is AI-generated or human-made."
)
if __name__ == "__main__":
demo.launch() |