Update app.py

app.py

@@ -1,10 +1,12 @@
 import spaces
 import gradio as gr
-from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassification
+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
 
 # Suppress warnings
 warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")
@@ -20,13 +22,24 @@ clf_1 = pipeline(model=model_1, task="image-classification", image_processor=ima
 
 # Load the second model
 model_2_path = "Heem2/AI-vs-Real-Image-Detection"
-clf_2 = pipeline("image-classification", model=model_2_path, device=device)
+clf_2 = pipeline("image-classification", model=model_2_path)
 
-# Define class names for both models
+# 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']  # Adjust if the second model has different classes
+class_names_2 = ['AI Image', 'Real Image']
+class_names_3 = ['AI', 'Real']
+class_names_4 = ['AI', 'Real']
+
+def softmax(vector):
+    e = math.exp(vector - vector.max())  # for numerical stability
+    return e / e.sum()
 
-@spaces.GPU(duration=30)
+@spaces.GPU(duration=10)
 def predict_image(img, confidence_threshold):
     # Ensure the image is a PIL Image
     if not isinstance(img, Image.Image):
@@ -81,10 +94,56 @@ def predict_image(img, confidence_threshold):
     except Exception as e:
         label_2 = f"Error: {str(e)}"
     
+    # Predict using the third model
+    try:
+        prediction_3 = pipe0(img_pil)
+        result_3 = {}
+        for idx, result in enumerate(prediction_3):
+            result_3[class_names_3[idx]] = float(result['score'])
+        
+        # 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
+    try:
+        prediction_4 = pipe1(img_pil)
+        result_4 = {}
+        for idx, result in enumerate(prediction_4):
+            result_4[class_names_4[idx]] = float(result['score'])
+        
+        # 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
+        "AI-vs-Real-Image-Detection": label_2,
+        "Organika/sdxl-detector": label_3,
+        "cmckinle/sdxl-flux-detector": label_4
     }
     
     return combined_results