Spaces:
Runtime error
Runtime error
| import spaces | |
| import gradio as gr | |
| from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification, AutoFeatureExtractor, AutoModelForImageClassification | |
| from torchvision import transforms | |
| import torch | |
| from PIL import Image | |
| import pandas as pd | |
| import warnings | |
| import math | |
| import numpy as np | |
| # 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) | |
| # Load additional models | |
| models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"] | |
| pipe0 = pipeline("image-classification", model=models[0]) | |
| pipe1 = pipeline("image-classification", model=models[1]) | |
| # Define class names for all models | |
| class_names_1 = ['artificial', 'real'] | |
| class_names_2 = ['AI Image', 'Real Image'] | |
| class_names_3 = ['AI', 'Real'] | |
| class_names_4 = ['AI', 'Real'] | |
| def softmax(vector): | |
| e = np.exp(vector - np.max(vector)) # for numerical stability | |
| return e / e.sum() | |
| 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)}" | |
| # Predict using the third model with softmax | |
| try: | |
| with torch.no_grad(): | |
| outputs = pipe0(img_pil) | |
| logits = outputs[0].logits if isinstance(outputs, list) else outputs.logits | |
| probabilities = softmax(logits.cpu().numpy()) | |
| result_3 = {class_names_3[idx]: float(probabilities[idx]) for idx in range(len(class_names_3))} | |
| # Ensure the result dictionary contains all class names | |
| for class_name in class_names_3: | |
| if class_name not in result_3: | |
| result_3[class_name] = 0.0 | |
| # Check if either class meets the confidence threshold | |
| if result_3['AI'] >= confidence_threshold: | |
| label_3 = f"Label: AI, Confidence: {result_3['AI']:.4f}" | |
| elif result_3['Real'] >= confidence_threshold: | |
| label_3 = f"Label: Real, Confidence: {result_3['Real']:.4f}" | |
| else: | |
| label_3 = "Uncertain Classification" | |
| except Exception as e: | |
| label_3 = f"Error: {str(e)}" | |
| # Predict using the fourth model with softmax | |
| try: | |
| with torch.no_grad(): | |
| outputs = pipe1(img_pil) | |
| logits = outputs[0].logits if isinstance(outputs, list) else outputs.logits | |
| probabilities = softmax(logits.cpu().numpy()) | |
| result_4 = {class_names_4[idx]: float(probabilities[idx]) for idx in range(len(class_names_4))} | |
| # Ensure the result dictionary contains all class names | |
| for class_name in class_names_4: | |
| if class_name not in result_4: | |
| result_4[class_name] = 0.0 | |
| # Check if either class meets the confidence threshold | |
| if result_4['AI'] >= confidence_threshold: | |
| label_4 = f"Label: AI, Confidence: {result_4['AI']:.4f}" | |
| elif result_4['Real'] >= confidence_threshold: | |
| label_4 = f"Label: Real, Confidence: {result_4['Real']:.4f}" | |
| else: | |
| label_4 = "Uncertain Classification" | |
| except Exception as e: | |
| label_4 = f"Error: {str(e)}" | |
| # Combine results | |
| combined_results = { | |
| "SwinV2": label_1, | |
| "AI-vs-Real-Image-Detection": label_2, | |
| "Organika/sdxl-detector": label_3, | |
| "cmckinle/sdxl-flux-detector": label_4 | |
| } | |
| 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") | |
| # Launch the interface | |
| iface = gr.Interface( | |
| fn=predict_image, | |
| inputs=[image, confidence_slider], | |
| outputs=label, | |
| title="AI Generated Classification" | |
| ) | |
| iface.launch() |