File size: 2,057 Bytes
d970d90 |
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 |
import gradio as gr
import torch
from PIL import Image as PILImage
from transformers import AutoImageProcessor, SiglipForImageClassification
# --- Model Loading ---
MODEL_IDENTIFIER = r"Ateeqq/ai-vs-human-image-detector"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Loading model: {MODEL_IDENTIFIER} on device: {device}")
processor = AutoImageProcessor.from_pretrained(MODEL_IDENTIFIER)
model = SiglipForImageClassification.from_pretrained(MODEL_IDENTIFIER)
model.to(device)
model.eval()
print("Model loaded successfully.")
# --- Prediction Function ---
def predict(image):
"""
Takes a PIL image, preprocesses it, and returns the prediction probabilities.
"""
if image is None:
return None
# Preprocess the image
inputs = processor(images=image, return_tensors="pt").to(device)
# Perform inference
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
# Get probabilities
probabilities = torch.softmax(logits, dim=-1)[0]
# Create a dictionary of labels and their scores
confidences = {model.config.id2label[i]: score.item() for i, score in enumerate(probabilities)}
return confidences
# --- Gradio Interface ---
# Define the Gradio interface components
image_input = gr.Image(type="pil", label="Upload an Image")
label_output = gr.Label(num_top_classes=2, label="Prediction")
# The title and description for the app
title = "AI vs Human Image Detector"
description = """
This Space uses the `Ateeqq/ai-vs-human-image-detector` model to classify an image as either AI-generated or Human-made.
Upload an image to see the prediction.
"""
article = "Model by [Ateeqq](https://huggingface.co/Ateeqq) | Gradio app created with AI"
# Launch the Gradio app
gr.Interface(
fn=predict,
inputs=image_input,
outputs=label_output,
title=title,
description=description,
article=article
).launch(share=True, server_name="0.0.0.0") |