Spaces:
Runtime error
Runtime error
import spaces | |
import gradio as gr | |
from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification | |
from torchvision import transforms | |
import torch | |
from PIL import Image | |
import warnings | |
# Suppress warnings | |
warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset") | |
# Ensure using GPU if available | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
# Load the first model and processor | |
image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy", use_fast=True) | |
model_1 = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy") | |
model_1 = model_1.to(device) | |
clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device) | |
# Load the second model | |
model_2_path = "Heem2/AI-vs-Real-Image-Detection" | |
clf_2 = pipeline("image-classification", model=model_2_path) | |
# Define class names for both models | |
class_names_1 = ['artificial', 'real'] | |
class_names_2 = ['AI Image', 'Real Image'] # Adjust if the second model has different classes | |
def predict_image(img, confidence_threshold): | |
# Ensure the image is a PIL Image | |
if not isinstance(img, Image.Image): | |
raise ValueError(f"Expected a PIL Image, but got {type(img)}") | |
# Convert the image to RGB if not already | |
if img.mode != 'RGB': | |
img_pil = img.convert('RGB') | |
else: | |
img_pil = img | |
# Resize the image | |
img_pil = transforms.Resize((256, 256))(img_pil) | |
# Predict using the first model | |
try: | |
prediction_1 = clf_1(img_pil) | |
result_1 = {pred['label']: pred['score'] for pred in prediction_1} | |
# Ensure the result dictionary contains all class names | |
for class_name in class_names_1: | |
if class_name not in result_1: | |
result_1[class_name] = 0.0 | |
# Check if either class meets the confidence threshold | |
if result_1['artificial'] >= confidence_threshold: | |
label_1 = f"Label: artificial, Confidence: {result_1['artificial']:.4f}" | |
elif result_1['real'] >= confidence_threshold: | |
label_1 = f"Label: real, Confidence: {result_1['real']:.4f}" | |
else: | |
label_1 = "Uncertain Classification" | |
except Exception as e: | |
label_1 = f"Error: {str(e)}" | |
# Predict using the second model | |
try: | |
prediction_2 = clf_2(img_pil) | |
result_2 = {pred['label']: pred['score'] for pred in prediction_2} | |
# Ensure the result dictionary contains all class names | |
for class_name in class_names_2: | |
if class_name not in result_2: | |
result_2[class_name] = 0.0 | |
# Check if either class meets the confidence threshold | |
if result_2['AI Image'] >= confidence_threshold: | |
label_2 = f"Label: AI Image, Confidence: {result_2['AI Image']:.4f}" | |
elif result_2['Real Image'] >= confidence_threshold: | |
label_2 = f"Label: Real Image, Confidence: {result_2['Real Image']:.4f}" | |
else: | |
label_2 = "Uncertain Classification" | |
except Exception as e: | |
label_2 = f"Error: {str(e)}" | |
# Combine results | |
combined_results = { | |
"SwinV2": label_1, | |
"AI-vs-Real-Image-Detection": label_2 | |
} | |
return combined_results | |
# Define the Gradio interface | |
image = gr.Image(label="Image to Analyze", sources=['upload'], type='pil') # Ensure the image type is PIL | |
confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold") | |
label = gr.JSON(label="Model Predictions") | |
gr.Interface( | |
fn=predict_image, | |
inputs=[image, confidence_slider], | |
outputs=label, | |
title="AI Generated Classification", | |
queue=True # Enable queuing to handle multiple predictions efficiently | |
).launch() |