Update app.py

app.py

@@ -7,9 +7,7 @@ import shutil
 import matplotlib.pyplot as plt
 from sklearn.metrics import accuracy_score, roc_auc_score, confusion_matrix, classification_report, roc_curve, auc, ConfusionMatrixDisplay
 from PIL import Image
-import uuid
 import tempfile
-import pandas as pd
 import numpy as np
 import urllib.request
 
@@ -21,19 +19,12 @@ class AIDetector:
         self.pipe = pipeline("image-classification", MODEL_NAME)
         self.feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
         self.model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
-        self.results = []
 
     @staticmethod
     def softmax(vector):
         e = np.exp(vector - np.max(vector))
         return e / e.sum()
 
-    def classify_image(self, image):
-        outputs = self.pipe(image)
-        results = {label: float(output['score']) for label, output in zip(LABELS, outputs)}
-        self.results.append(results)
-        return results
-
     def predict(self, image):
         inputs = self.feature_extractor(image, return_tensors="pt")
         with torch.no_grad():
@@ -45,19 +36,9 @@ class AIDetector:
         label = LABELS[prediction]
         
         results = {label: float(prob) for label, prob in zip(LABELS, probabilities[0])}
-        self.results.append(results)
         
         return label, results
 
-    def get_total_probability(self):
-        if not self.results:
-            return None
-        avg_real_prob = sum(result["Real"] for result in self.results) / len(self.results)
-        return {"Real": f"{avg_real_prob:.4f}", "AI": f"{1 - avg_real_prob:.4f}"}
-
-    def clear_results(self):
-        self.results.clear()
-
 def process_zip(zip_file):
     temp_dir = tempfile.mkdtemp()
     with zipfile.ZipFile(zip_file.name, 'r') as z:
@@ -140,13 +121,10 @@ def create_gradio_interface():
                     message = gr.HTML()
 
                 with gr.Group():
-                    with gr.Row():
-                        final_prob = gr.Label(label="Final Probability")
-                    with gr.Row():
-                        with gr.Box():
-                            gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{MODEL_NAME}'>{MODEL_NAME}</a></b>""")
-                            output_html = gr.HTML()
-                            output_label = gr.Label(label="Output")
+                    with gr.Box():
+                        gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{MODEL_NAME}'>{MODEL_NAME}</a></b>""")
+                        output_html = gr.HTML()
+                        output_label = gr.Label(label="Output")
 
             with gr.Tab("Batch Image Processing"):
                 zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
@@ -160,16 +138,10 @@ def create_gradio_interface():
                     output_plots = gr.Plot(label="Confusion Matrix and ROC Curve")
 
         load_btn.click(load_url, in_url, [inp, message])
-        btn.click(detector.clear_results, None, final_prob, show_progress=False)
         btn.click(
             lambda img: detector.predict(img),
             inp,
             [output_html, output_label]
-        ).then(
-            lambda: detector.get_total_probability(),
-            None,
-            final_prob,
-            show_progress=False
         )
 
         batch_btn.click(