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 numpy as np import io import logging from utils.utils import softmax, augment_image, convert_pil_to_bytes # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Ensure using GPU if available device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # Model paths and class names MODEL_PATHS = { "model_1": "haywoodsloan/ai-image-detector-deploy", "model_2": "Heem2/AI-vs-Real-Image-Detection", "model_3": "Organika/sdxl-detector", "model_4": "cmckinle/sdxl-flux-detector", "model_5": "prithivMLmods/Deep-Fake-Detector-v2-Model", "model_5b": "prithivMLmods/Deepfake-Detection-Exp-02-22" } CLASS_NAMES = { "model_1": ['artificial', 'real'], "model_2": ['AI Image', 'Real Image'], "model_3": ['AI', 'Real'], "model_4": ['AI', 'Real'], "model_5": ['Realism', 'Deepfake'], "model_5b": ['Real', 'Deepfake'] } # Load models and processors def load_models(): image_processor_1 = AutoImageProcessor.from_pretrained(MODEL_PATHS["model_1"], use_fast=True) model_1 = Swinv2ForImageClassification.from_pretrained(MODEL_PATHS["model_1"]) model_1 = model_1.to(device) clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device) clf_2 = pipeline("image-classification", model=MODEL_PATHS["model_2"], device=device) feature_extractor_3 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_3"], device=device) model_3 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_3"]).to(device) feature_extractor_4 = AutoFeatureExtractor.from_pretrained(MODEL_PATHS["model_4"], device=device) model_4 = AutoModelForImageClassification.from_pretrained(MODEL_PATHS["model_4"]).to(device) clf_5 = pipeline("image-classification", model=MODEL_PATHS["model_5"], device=device) clf_5b = pipeline("image-classification", model=MODEL_PATHS["model_5b"], device=device) return clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b clf_1, clf_2, feature_extractor_3, model_3, feature_extractor_4, model_4, clf_5, clf_5b = load_models() @spaces.GPU(duration=10) def predict_with_model(img_pil, clf, class_names, confidence_threshold, model_name, model_id): try: prediction = clf(img_pil) result = {pred['label']: pred['score'] for pred in prediction} result_output = [model_id, model_name, result.get(class_names[1], 0.0), result.get(class_names[0], 0.0)] logger.info(result_output) for class_name in class_names: if class_name not in result: result[class_name] = 0.0 if result[class_names[0]] >= confidence_threshold: label = f"AI, Confidence: {result[class_names[0]]:.4f}" result_output.append('AI') elif result[class_names[1]] >= confidence_threshold: label = f"Real, Confidence: {result[class_names[1]]:.4f}" result_output.append('REAL') else: label = "Uncertain Classification" result_output.append('UNCERTAIN') except Exception as e: label = f"Error: {str(e)}" return label, result_output @spaces.GPU(duration=10) # app.py def predict_image(img, confidence_threshold): if not isinstance(img, Image.Image): raise ValueError(f"Expected a PIL Image, but got {type(img)}") if img.mode != 'RGB': img_pil = img.convert('RGB') else: img_pil = img img_pil = transforms.Resize((256, 256))(img_pil) img_pilvits = transforms.Resize((224, 224))(img_pil) label_1, result_1output = predict_with_model(img_pil, clf_1, CLASS_NAMES["model_1"], confidence_threshold, "SwinV2-base", 1) label_2, result_2output = predict_with_model(img_pilvits, clf_2, CLASS_NAMES["model_2"], confidence_threshold, "ViT-base Classifier", 2) label_3, result_3output = predict_with_model(img_pil, model_3, CLASS_NAMES["model_3"], confidence_threshold, "SDXL-Trained", 3) label_4, result_4output = predict_with_model(img_pil, model_4, CLASS_NAMES["model_4"], confidence_threshold, "SDXL + FLUX", 4) label_5, result_5output = predict_with_model(img_pilvits, clf_5, CLASS_NAMES["model_5"], confidence_threshold, "ViT-base Newcomer", 5) label_5b, result_5boutput = predict_with_model(img_pilvits, clf_5b, CLASS_NAMES["model_5b"], confidence_threshold, "ViT-base Newcomer", 6) combined_results = { "SwinV2/detect": label_1, "ViT/AI-vs-Real": label_2, "Swin/SDXL": label_3, "Swin/SDXL-FLUX": label_4, "prithivMLmods": label_5, "prithivMLmods-2-22": label_5b } combined_outputs = [result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput] return img_pil, combined_outputs # Define a function to generate the HTML content # Define a function to generate the HTML content def generate_results_html(results): def get_header_color(label): if label == 'AI': return 'bg-red-500 text-red-700', 'bg-red-400', 'bg-red-100', 'bg-red-700 text-red-700', 'bg-red-200' elif label == 'REAL': return 'bg-green-500 text-green-700', 'bg-green-400', 'bg-green-100', 'bg-green-700 text-green-700', 'bg-green-200' elif label == 'UNCERTAIN': return 'bg-yellow-500 text-yellow-700 bg-yellow-100', 'bg-yellow-400', 'bg-yellow-100', 'bg-yellow-700 text-yellow-700', 'bg-yellow-200' elif label == 'MAINTENANCE': return 'bg-blue-500 text-blue-700', 'bg-blue-400', 'bg-blue-100', 'bg-blue-700 text-blue-700', 'bg-blue-200' else: return 'bg-gray-300 text-gray-700', 'bg-gray-400', 'bg-gray-100', 'bg-gray-700 text-gray-700', 'bg-gray-200' html_content = f"""
MODEL 1: {'' if results[0][-1] == 'REAL' else ''}

{results[0][-1]}

Conf: {results[0][2]:.4f}

Conf: {results[0][3]:.4f}

SwinV2 Based

Real: {results[0][2]:.4f}, AI: {results[0][3]:.4f}
@haywoodsloan / more info
MODEL 2: {'' if results[1][-1] == 'REAL' else ''}

{results[1][-1]}

Conf: {results[1][2]:.4f}

Conf: {results[1][3]:.4f}

ViT Based

Real: {results[1][2]:.4f}, AI: {results[1][3]:.4f}
@Heem2 / more info
MODEL 3: {'' if results[2][-1] == 'REAL' else ''}

{results[2][-1]}

Conf: {results[2][2]:.4f}

Conf: {results[2][3]:.4f}

SDXL Dataset

Real: {results[2][2]:.4f}, AI: {results[2][3]:.4f}
@Organika / more info
MODEL 4: {'' if results[3][-1] == 'REAL' else ''}

{results[3][-1]}

Conf: {results[3][2]:.4f}

Conf: {results[3][3]:.4f}

SDXL + FLUX

Real: {results[3][2]:.4f}, AI: {results[3][3]:.4f}
@cmckinle / more info
MODEL 5: {'' if results[4][-1] == 'REAL' else ''}

{results[4][-1]}

Conf: {results[4][2]:.4f}

Conf: {results[4][3]:.4f}

Vits Model

Real: {results[4][2]:.4f}, AI: {results[4][3]:.4f}
Real: {results[5][2]:.4f}, AI: {results[5][3]:.4f}
@prithivMLmods / more info
""" return html_content # Modify the predict_image function to return the HTML content def predict_image_with_html(img, confidence_threshold): img_pil, results = predict_image(img, confidence_threshold) html_content = generate_results_html(results) return img_pil, html_content # Define the Gradio interface with gr.Blocks() as iface: gr.Markdown("# AI Generated Image / Deepfake Detection Models Evaluation") with gr.Row(): with gr.Column(scale=1): image_input = gr.Image(label="Upload Image to Analyze", sources=['upload'], type='pil') with gr.Accordion("Settings", open=False, elem_id="settings_accordion"): confidence_slider = gr.Slider(0.0, 1.0, value=0.5, step=0.01, label="Confidence Threshold") inputs = [image_input, confidence_slider] with gr.Column(scale=2): with gr.Accordion("Project OpenSight - Model Evaluations & Playground", open=True, elem_id="project_accordion"): gr.Markdown("## OpenSight is a SOTA gen. image detection model, in pre-release prep.\n\nThis HF Space is a temporary home for us and the public to evaluate the shortcomings of current open source models.\n\n<-- Feel free to play around by starting with an image as we prepare our formal announcement.") image_output = gr.Image(label="Processed Image", visible=False) # Custom HTML component to display results in 5 columns results_html = gr.HTML(label="Model Predictions") outputs = [image_output, results_html] # gr.Button("Predict").click(fn=predict_image_with_html, inputs=inputs, outputs=outputs) predict_button = gr.Button("Predict") predict_button.click( fn=predict_image_with_html, inputs=inputs, outputs=outputs ) predict_button.click( fn=None, js="() => {document.getElementById('project_accordion').open = false;}", # Close the project accordion inputs=[], outputs=[] ) # Launch the interface iface.launch()