Enhance prediction output and result display in app.py

- Modify prediction methods to generate structured output lists for each model
- Add model-specific output tracking with confidence scores and classification labels
- Update HTML results display to include model badges
- Adjust Gradio interface layout for better visualization
- Improve error handling and logging in prediction functions

.gitignore

@@ -0,0 +1 @@
+.env

app.py

@@ -67,7 +67,8 @@ def predict_image(img, confidence_threshold):
     try:
         prediction_1 = clf_1(img_pil)
         result_1 = {pred['label']: pred['score'] for pred in prediction_1}
-        print(result_1)
+        result_1output = [1, result_1['real'], result_1['artificial']]
+        print(result_1output)
         # Ensure the result dictionary contains all class names
         for class_name in class_names_1:
             if class_name not in result_1:
@@ -75,18 +76,23 @@ def predict_image(img, confidence_threshold):
         # Check if either class meets the confidence threshold
         if result_1['artificial'] >= confidence_threshold:
             label_1 = f"AI, Confidence: {result_1['artificial']:.4f}"
+            result_1output += ['AI']
         elif result_1['real'] >= confidence_threshold:
             label_1 = f"Real, Confidence: {result_1['real']:.4f}"
+            result_1output += ['REAL']
         else:
             label_1 = "Uncertain Classification"
+            result_1output += ['UNCERTAIN']
+
     except Exception as e:
         label_1 = f"Error: {str(e)}"
-    
+    print(result_1output)
     # Predict using the second model
     try:
         prediction_2 = clf_2(img_pil)
         result_2 = {pred['label']: pred['score'] for pred in prediction_2}
-        print(result_2)
+        result_2output = [2, result_2['Real Image'], result_2['AI Image']]
+        print(result_2output)
         # Ensure the result dictionary contains all class names
         for class_name in class_names_2:
             if class_name not in result_2:
@@ -94,10 +100,13 @@ def predict_image(img, confidence_threshold):
         # Check if either class meets the confidence threshold
         if result_2['AI Image'] >= confidence_threshold:
             label_2 = f"AI, Confidence: {result_2['AI Image']:.4f}"
+            result_2output += ['AI']
         elif result_2['Real Image'] >= confidence_threshold:
             label_2 = f"Real, Confidence: {result_2['Real Image']:.4f}"
+            result_2output += ['REAL']
         else:
             label_2 = "Uncertain Classification"
+            result_2output += ['UNCERTAIN']
     except Exception as e:
         label_2 = f"Error: {str(e)}"
     
@@ -109,10 +118,11 @@ def predict_image(img, confidence_threshold):
             logits_3 = outputs_3.logits
             probabilities_3 = softmax(logits_3.cpu().numpy()[0])
         result_3 = {
-            labels_3[0]: float(probabilities_3[0]),  # AI
-            labels_3[1]: float(probabilities_3[1])   # Real
+            labels_3[1]: float(probabilities_3[1]),   # Real
+            labels_3[0]: float(probabilities_3[0])  # AI
         }
-        print(result_3)
+        result_3output = [3, float(probabilities_3[1]), float(probabilities_3[0])]
+        print(result_3output)
         # Ensure the result dictionary contains all class names
         for class_name in labels_3:
             if class_name not in result_3:
@@ -120,10 +130,13 @@ def predict_image(img, confidence_threshold):
         # Check if either class meets the confidence threshold
         if result_3['AI'] >= confidence_threshold:
             label_3 = f"AI, Confidence: {result_3['AI']:.4f}"
+            result_3output += ['AI']
         elif result_3['Real'] >= confidence_threshold:
             label_3 = f"Real, Confidence: {result_3['Real']:.4f}"
+            result_3output += ['REAL']
         else:
             label_3 = "Uncertain Classification"
+            result_3output += ['UNCERTAIN']
     except Exception as e:
         label_3 = f"Error: {str(e)}"
     
@@ -135,9 +148,10 @@ def predict_image(img, confidence_threshold):
             logits_4 = outputs_4.logits
             probabilities_4 = softmax(logits_4.cpu().numpy()[0])
         result_4 = {
-            labels_4[0]: float(probabilities_4[0]),  # AI
-            labels_4[1]: float(probabilities_4[1])   # Real
+            labels_4[1]: float(probabilities_4[1]),   # Real
+            labels_4[0]: float(probabilities_4[0])  # AI
         }
+        result_4output = [4, float(probabilities_4[1]), float(probabilities_4[0])]
         print(result_4)
         # Ensure the result dictionary contains all class names
         for class_name in labels_4:
@@ -146,19 +160,27 @@ def predict_image(img, confidence_threshold):
         # Check if either class meets the confidence threshold
         if result_4['AI'] >= confidence_threshold:
             label_4 = f"AI, Confidence: {result_4['AI']:.4f}"
+            result_4output += ['AI']
         elif result_4['Real'] >= confidence_threshold:
             label_4 = f"Real, Confidence: {result_4['Real']:.4f}"
+            result_4output += ['REAL']
         else:
             label_4 = "Uncertain Classification"
+            result_4output += ['UNCERTAIN']
     except Exception as e:
         label_4 = f"Error: {str(e)}"
     
     try:
+        result_5output = [5, 0.0, 0.0, 'MAINTENANCE']
         img_bytes = convert_pil_to_bytes(img_pil)
-        response5_raw = call_inference(img_bytes)
-        response5 = response5_raw.json()
+        # 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)}"
     
@@ -170,32 +192,34 @@ def predict_image(img, confidence_threshold):
         "Swin/SDXL-FLUX": label_4,
         "GOAT": label_5
     }
-    return img_pil, combined_results
+    combined_outputs = [ result_1output, result_2output, result_3output, result_4output, result_5output ]
+    return img_pil, combined_outputs
 
 # Define a function to generate the HTML content
 def generate_results_html(results):
+    print(results)
     html_content = f"""
     <link href="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/css/bootstrap.min.css" rel="stylesheet">
     <div class="container">
         <div class="row mt-4">
             <div class="col">
-                <h5>SwinV2/detect</h5>
+                <h5>SwinV2/detect <span class="badge badge-secondary">M1</span></h5>
                 <p>{results.get("SwinV2/detect", "N/A")}</p>
             </div>
             <div class="col">
-                <h5>ViT/AI-vs-Real</h5>
+                <h5>ViT/AI-vs-Real <span class="badge badge-secondary">M2</span></h5>
                 <p>{results.get("ViT/AI-vs-Real", "N/A")}</p>
             </div>
             <div class="col">
-                <h5>Swin/SDXL</h5>
+                <h5>Swin/SDXL <span class="badge badge-secondary">M3</span></h5>
                 <p>{results.get("Swin/SDXL", "N/A")}</p>
             </div>
             <div class="col">
-                <h5>Swin/SDXL-FLUX</h5>
+                <h5>Swin/SDXL-FLUX <span class="badge badge-secondary">M4</span></h5>
                 <p>{results.get("Swin/SDXL-FLUX", "N/A")}</p>
             </div>
             <div class="col">
-                <h5>GOAT</h5>
+                <h5>GOAT <span class="badge badge-secondary">M5</span></h5>
                 <p>{results.get("GOAT", "N/A")}</p>
             </div>
         </div>
@@ -214,11 +238,11 @@ with gr.Blocks() as iface:
     gr.Markdown("# AI Generated Image Classification")
     
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=2):
             image_input = gr.Image(label="Upload Image to Analyze", sources=['upload'], type='pil')
             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():
+        with gr.Column(scale=3):
             image_output = gr.Image(label="Processed Image")
             # Custom HTML component to display results in 5 columns
             results_html = gr.HTML(label="Model Predictions")