Update app.py

app.py

@@ -36,12 +36,16 @@ model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)
 model_5_path = "prithivMLmods/Deep-Fake-Detector-v2-Model"
 clf_5 = pipeline("image-classification", model=model_5_path, device=device)
 
+model_5b_path = "prithivMLmods/Deepfake-Detection-Exp-02-22"
+clf_5b = pipeline("image-classification", model=model_5b_path, device=device)
+
 # Define class names for all models
 class_names_1 = ['artificial', 'real']
 class_names_2 = ['AI Image', 'Real Image']
 labels_3 = ['AI', 'Real']
 labels_4 = ['AI', 'Real']
 class_names_5 = ['Realism', 'Deepfake']
+class_names_5b = ['Real', 'Deepfake']
 
 def softmax(vector):
     e = np.exp(vector - np.max(vector))  # for numerical stability
@@ -223,6 +227,30 @@ def predict_image(img, confidence_threshold):
         
     print(result_5output)
     
+    try:
+        prediction_5b = clf_5b(img_pilvits)
+        result_5b = {pred['label']: pred['score'] for pred in prediction_5b}
+        result_5boutput = [6, 'ViT-base Newcomer', result_5b['Real'], result_5b['Deepfake']]
+        
+        # Ensure the result dictionary contains all class names
+        for class_name in class_names_5b:
+            if class_name not in result_5b:
+                result_5b[class_name] = 0.0
+        # Check if either class meets the confidence threshold
+        if result_5b['Deepfake'] >= confidence_threshold:
+            label_5b = f"AI, Confidence: {result_5b['Deepfake']:.4f}"
+            result_5boutput += ['AI']
+        elif result_5b['Real Image'] >= confidence_threshold:
+            label_5b = f"Real, Confidence: {result_5b['Real']:.4f}"
+            result_5boutput += ['REAL']
+        else:
+            label_5b = "Uncertain Classification"
+            result_5boutput += ['UNCERTAIN']
+    except Exception as e:
+        label_5b = f"Error: {str(e)}"
+        
+    print(result_5boutput)
+    
     
     # try:
     #     result_5output = [5, 'TBA', 0.0, 0.0, 'MAINTENANCE']
@@ -245,11 +273,12 @@ def predict_image(img, confidence_threshold):
         "ViT/AI-vs-Real": label_2,
         "Swin/SDXL": label_3,
         "Swin/SDXL-FLUX": label_4,
-        "prithivMLmods": label_5
+        "prithivMLmods": label_5,
+        "prithivMLmods-2-22": label_5b 
     }
         # Generate HTML content
     
-    combined_outputs = [ result_1output, result_2output, result_3output, result_4output, result_5output ]
+    combined_outputs = [ result_1output, result_2output, result_3output, result_4output, result_5output, result_5boutput ]
     # html_content = generate_results_html(combined_outputs)
 
     return img_pil, combined_outputs
@@ -488,6 +517,7 @@ def generate_results_html(results):
                     <h4 class="mt-4 text-sm font-semibold tracking-wide">Vits Model</h4>
                     
                     <div class="text-xs font-mono">Real: {results[4][2]:.4f}, AI: {results[4][3]:.4f}</div>
+                    <div class="text-xs font-mono">Real: {results[5][2]:.4f}, AI: {results[5][3]:.4f}</div>
                     
                     <a class="mt-2 text-xs tracking-wide">@prithivMLmods / more info</a>
                 </div>