Add new ViT-based deep fake detection model

- Integrate prithivMLmods/Deep-Fake-Detector-v2-Model as the fifth classification model
- Update prediction logic to handle the new model's classification
- Modify README to include the new model in preload_from_hub
- Adjust result tile styling and labels for the new model
- Resize input images to 224x224 for ViT models

README.md

@@ -1,10 +1,10 @@
 ---
-title: Swin-V2-Classify
+title: OpenSight
 emoji: 🔥
 colorFrom: indigo
 colorTo: purple
 sdk: gradio
-sdk_version: 4.44.1
+sdk_version: 5.15.0
 app_file: app.py
 pinned: false
 preload_from_hub:
@@ -12,6 +12,7 @@ preload_from_hub:
   - haywoodsloan/ai-image-detector-deploy
   - cmckinle/sdxl-flux-detector
   - Organika/sdxl-detector
+  - prithivMLmods/Deep-Fake-Detector-v2-Model
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -32,11 +32,16 @@ model_3 = AutoModelForImageClassification.from_pretrained(models[0]).to(device)
 feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1], device=device)
 model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)
 
+# Load the second model
+model_5_path = "prithivMLmods/Deep-Fake-Detector-v2-Model"
+clf_5 = pipeline("image-classification", model=model_5_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']
 
 def softmax(vector):
     e = np.exp(vector - np.max(vector))  # for numerical stability
@@ -84,6 +89,8 @@ def predict_image(img, confidence_threshold):
 
     # Resize the image
     img_pil = transforms.Resize((256, 256))(img_pil)
+    # Size 224 for vits models
+    img_pilvits = transforms.Resize((224, 224))(img_pilvits)
     
     # Predict using the first model
     try:
@@ -111,7 +118,7 @@ def predict_image(img, confidence_threshold):
     print(result_1output)
     # Predict using the second model
     try:
-        prediction_2 = clf_2(img_pil)
+        prediction_2 = clf_2(img_pilvits)
         result_2 = {pred['label']: pred['score'] for pred in prediction_2}
         result_2output = [2, 'ViT-base Classifer', result_2['Real Image'], result_2['AI Image']]
         print(result_2output)
@@ -193,19 +200,42 @@ def predict_image(img, confidence_threshold):
         label_4 = f"Error: {str(e)}"
     
     try:
-        result_5output = [5, 'TBA', 0.0, 0.0, 'MAINTENANCE']
-        img_bytes = convert_pil_to_bytes(img_pil)
-        # print(img)
-        # print(img_bytes)
-        response5_raw = call_inference(img)
-        print(response5_raw)
-        response5 = response5_raw
-        print(response5)
-        label_5 = f"Result: {response5}"
-        
+        prediction_5 = clf_5(img_pilvits)
+        result_5 = {pred['label']: pred['score'] for pred in prediction_2}
+        result_5output = [5, 'ViT-base Newcomer', result_5['Realism'], result_5['Deepfake']]
+        print(result_2output)
+        # Ensure the result dictionary contains all class names
+        for class_name in class_names_5:
+            if class_name not in result_5:
+                result_5[class_name] = 0.0
+        # Check if either class meets the confidence threshold
+        if result_5['AI Image'] >= confidence_threshold:
+            label_5 = f"AI, Confidence: {result_5['Deepfake']:.4f}"
+            result_5output += ['AI']
+        elif result_5['Real Image'] >= confidence_threshold:
+            label_5 = f"Real, Confidence: {result_5['Realism']:.4f}"
+            result_5output += ['REAL']
+        else:
+            label_5 = "Uncertain Classification"
+            result_5output += ['UNCERTAIN']
     except Exception as e:
         label_5 = f"Error: {str(e)}"
     
+    
+    # try:
+    #     result_5output = [5, 'TBA', 0.0, 0.0, 'MAINTENANCE']
+    #     img_bytes = convert_pil_to_bytes(img_pil)
+    #     # print(img)
+    #     # print(img_bytes)
+    #     response5_raw = call_inference(img)
+    #     print(response5_raw)
+    #     response5 = response5_raw
+    #     print(response5)
+    #     label_5 = f"Result: {response5}"
+        
+    # except Exception as e:
+    #     label_5 = f"Error: {str(e)}"
+    
 
     # Combine results
     combined_results = {
@@ -213,7 +243,7 @@ def predict_image(img, confidence_threshold):
         "ViT/AI-vs-Real": label_2,
         "Swin/SDXL": label_3,
         "Swin/SDXL-FLUX": label_4,
-        "GOAT": label_5
+        "prithivMLmods": label_5
     }
         # Generate HTML content
     
@@ -285,7 +315,7 @@ def generate_results_html(results):
                     
                     <div class="text-xs font-mono">Real: {results[0][2]:.4f}, AI: {results[0][3]:.4f}</div>
                     
-                    <a class="mt-2 text-[0.66rem] tracking-wide">@haywoodsloan / more info</a>
+                    <a class="mt-2 text-xs tracking-wide">@haywoodsloan / more info</a>
                 </div>
             </div>
             <!-- Tile 2: ViT/AI-vs-Real -->
@@ -328,7 +358,7 @@ def generate_results_html(results):
                     
                     <div class="text-xs font-mono">Real: {results[1][2]:.4f}, AI: {results[1][3]:.4f}</div>
                     
-                    <a class="mt-2 text-[0.66rem] tracking-wide">@Heem2 / more info</a>
+                    <a class="mt-2 text-xs tracking-wide">@Heem2 / more info</a>
                 </div>
             </div>
             <!-- Tile 3: Swin/SDXL -->
@@ -367,11 +397,11 @@ def generate_results_html(results):
                     </div>
                 </div>
                 <div class="flex flex-col items-start">
-                    <h4 class="mt-4 text-sm font-semibold tracking-wide">Swin Based</h4>
+                    <h4 class="mt-4 text-sm font-semibold tracking-wide">SDXL Dataset</h4>
                     
                     <div class="text-xs font-mono">Real: {results[2][2]:.4f}, AI: {results[2][3]:.4f}</div>
                     
-                    <a class="mt-2 text-[0.66rem] tracking-wide">@Organika / more info</a>
+                    <a class="mt-2 text-xs tracking-wide">@Organika / more info</a>
                 </div>
             </div>
             <!-- Tile 4: Swin/SDXL-FLUX -->
@@ -410,11 +440,11 @@ def generate_results_html(results):
                     </div>
                 </div>
                 <div class="flex flex-col items-start">
-                    <h4 class="mt-4 text-sm font-semibold tracking-wide">Swin Based</h4>
+                    <h4 class="mt-4 text-sm font-semibold tracking-wide">SDXL + FLUX</h4>
                     
                     <div class="text-xs font-mono">Real: {results[3][2]:.4f}, AI: {results[3][3]:.4f}</div>
                     
-                    <a class="mt-2 text-[0.66rem] tracking-wide">@cmckinle / more info</a>
+                    <a class="mt-2 text-xs tracking-wide">@cmckinle / more info</a>
                 </div>
             </div>
             <!-- Tile 5: GOAT -->
@@ -453,11 +483,11 @@ def generate_results_html(results):
                     </div>
                 </div>
                 <div class="flex flex-col items-start">
-                    <h4 class="mt-4 text-sm font-semibold tracking-wide">GOAT Model</h4>
+                    <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>
                     
-                    <a class="mt-2 text-[0.66rem] tracking-wide">@GOAT / more info</a>
+                    <a class="mt-2 text-xs tracking-wide">@prithivMLmods / more info</a>
                 </div>
             </div>
         </div>