Update app.py

app.py

@@ -5,78 +5,69 @@ import os
 import zipfile
 import shutil
 import matplotlib.pyplot as plt
-from sklearn.metrics import accuracy_score, roc_auc_score, confusion_matrix, classification_report, roc_curve, auc
+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
-from numpy import exp
 import numpy as np
-from sklearn.metrics import ConfusionMatrixDisplay
 import urllib.request
 
-# Define model
-model = "cmckinle/sdxl-flux-detector"
-pipe = pipeline("image-classification", model)
-
-fin_sum = []
-uid = uuid.uuid4()
-
-# Softmax function
-def softmax(vector):
-    e = exp(vector - vector.max())  # for numerical stability
-    return e / e.sum()
-
-# Single image classification function
-def image_classifier(image):
-    labels = ["AI", "Real"]
-    outputs = pipe(image)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])
-    fin_sum.append(results)
-    return results
-
-def aiornot(image):
-    labels = ["AI", "Real"]
-    feature_extractor = AutoFeatureExtractor.from_pretrained(model)
-    model_cls = AutoModelForImageClassification.from_pretrained(model)
-    input = feature_extractor(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model_cls(**input)
-        logits = outputs.logits
-        probability = softmax(logits)
-        px = pd.DataFrame(probability.numpy())
-    prediction = logits.argmax(-1).item()
-    label = labels[prediction]
-
-    html_out = f"""
-    <h1>This image is likely: {label}</h1><br><h3>
-    Probabilities:<br>
-    Real: {float(px[1][0]):.4f}<br>
-    AI: {float(px[0][0]):.4f}"""
-    
-    results = {
-        "Real": float(px[1][0]),
-        "AI": float(px[0][0])
-    }
-    fin_sum.append(results)
-    return gr.HTML.update(html_out), results
-
-# Function to extract images from zip
-def extract_zip(zip_file):
+MODEL_NAME = "cmckinle/sdxl-flux-detector"
+LABELS = ["AI", "Real"]
+
+class AIDetector:
+    def __init__(self):
+        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():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probabilities = self.softmax(logits.numpy())
+        
+        prediction = logits.argmax(-1).item()
+        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, 'r') as z:
+    with zipfile.ZipFile(zip_file.name, 'r') as z:
         z.extractall(temp_dir)
-    return temp_dir
-
-# Function to classify images in a folder
-def classify_images(image_dir):
-    images = []
-    labels = []
-    preds = []
+    
+    labels, preds, images = [], [], []
+    detector = AIDetector()
+    
     for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
-        folder_path = os.path.join(image_dir, folder_name)
+        folder_path = os.path.join(temp_dir, folder_name)
         if not os.path.exists(folder_path):
             print(f"Folder not found: {folder_path}")
             continue
@@ -84,19 +75,18 @@ def classify_images(image_dir):
             img_path = os.path.join(folder_path, img_name)
             try:
                 img = Image.open(img_path).convert("RGB")
-                pred = pipe(img)
-                pred_label = 0 if pred[0]['label'] == 'AI' else 1
-
+                _, prediction = detector.predict(img)
+                pred_label = 0 if prediction["AI"] > prediction["Real"] else 1
+                
                 preds.append(pred_label)
                 labels.append(ground_truth_label)
                 images.append(img_name)
             except Exception as e:
                 print(f"Error processing image {img_name}: {e}")
     
-    print(f"Processed {len(images)} images")
-    return labels, preds, images
+    shutil.rmtree(temp_dir)
+    return evaluate_model(labels, preds)
 
-# Function to generate evaluation metrics
 def evaluate_model(labels, preds):
     cm = confusion_matrix(labels, preds)
     accuracy = accuracy_score(labels, preds)
@@ -105,105 +95,91 @@ def evaluate_model(labels, preds):
     fpr, tpr, _ = roc_curve(labels, preds)
     roc_auc = auc(fpr, tpr)
 
-    fig, ax = plt.subplots()
-    disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=["AI", "Real"])
-    disp.plot(cmap=plt.cm.Blues, ax=ax)
-    plt.close(fig)
-
-    fig_roc, ax_roc = plt.subplots()
-    ax_roc.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
-    ax_roc.plot([0, 1], [0, 1], color='gray', linestyle='--')
-    ax_roc.set_xlim([0.0, 1.0])
-    ax_roc.set_ylim([0.0, 1.05])
-    ax_roc.set_xlabel('False Positive Rate')
-    ax_roc.set_ylabel('True Positive Rate')
-    ax_roc.set_title('Receiver Operating Characteristic (ROC) Curve')
-    ax_roc.legend(loc="lower right")
-    plt.close(fig_roc)
-
-    return accuracy, roc_score, report, fig, fig_roc
-
-# Batch processing
-def process_zip(zip_file):
-    extracted_dir = extract_zip(zip_file.name)
-    labels, preds, images = classify_images(extracted_dir)
-    accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds)
-    shutil.rmtree(extracted_dir)  # Clean up extracted files
-    return accuracy, roc_score, report, cm_fig, roc_fig
+    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
+    
+    ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=LABELS).plot(cmap=plt.cm.Blues, ax=ax1)
+    ax1.set_title("Confusion Matrix")
+    
+    ax2.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
+    ax2.plot([0, 1], [0, 1], color='gray', linestyle='--')
+    ax2.set_xlim([0.0, 1.0])
+    ax2.set_ylim([0.0, 1.05])
+    ax2.set_xlabel('False Positive Rate')
+    ax2.set_ylabel('True Positive Rate')
+    ax2.set_title('ROC Curve')
+    ax2.legend(loc="lower right")
+    
+    plt.tight_layout()
+
+    return accuracy, roc_score, report, fig
 
-# Single image section
 def load_url(url):
     try:
-        urllib.request.urlretrieve(f'{url}', f"{uid}tmp_im.png")
-        image = Image.open(f"{uid}tmp_im.png")
-        mes = "Image Loaded"
+        urllib.request.urlretrieve(url, "temp_image.png")
+        image = Image.open("temp_image.png")
+        message = "Image Loaded"
     except Exception as e:
         image = None
-        mes = f"Image not Found<br>Error: {e}"
-    return image, mes
-
-def tot_prob():
-    try:
-        fin_out = sum([result["Real"] for result in fin_sum]) / len(fin_sum)
-        fin_sub = 1 - fin_out
-        out = {
-            "Real": f"{fin_out:.4f}",
-            "AI": f"{fin_sub:.4f}"
-        }
-        return out
-    except Exception as e:
-        print(e)
-        return None
-
-def fin_clear():
-    fin_sum.clear()
-    return None
-
-# Set up Gradio app
-with gr.Blocks() as app:
-    gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)</h4></h1></center>""")
-
-    with gr.Tabs():
-        # Tab for single image detection
-        with gr.Tab("Single Image Detection"):
-            with gr.Column():
-                inp = gr.Image(type='pil')
-                in_url = gr.Textbox(label="Image URL")
-                with gr.Row():
-                    load_btn = gr.Button("Load URL")
-                    btn = gr.Button("Detect AI")
-                mes = gr.HTML("""""")
-
-            with gr.Group():
-                with gr.Row():
-                    fin = gr.Label(label="Final Probability")
-                with gr.Row():
-                    with gr.Box():
-                        gr.HTML(f"""<b>Testing on Model: <a href='https://huggingface.co/{model}'>{model}</a></b>""")
-                        outp = gr.HTML("""""")
-                        n_out = gr.Label(label="Output")
-
-            btn.click(fin_clear, None, fin, show_progress=False)
-            load_btn.click(load_url, in_url, [inp, mes])
-
-            btn.click(aiornot, [inp], [outp, n_out]).then(
-                tot_prob, None, fin, show_progress=False)
-
-        # Tab for batch processing
-        with gr.Tab("Batch Image Processing"):
-            zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
-            batch_btn = gr.Button("Process Batch")
-
-            with gr.Group():
-                gr.Markdown(f"### Results for {model}")
-                output_acc = gr.Label(label="Accuracy")
-                output_roc = gr.Label(label="ROC Score")
-                output_report = gr.Textbox(label="Classification Report", lines=10)
-                output_cm = gr.Plot(label="Confusion Matrix")
-                output_roc_plot = gr.Plot(label="ROC Curve")
-
-    # Connect batch processing
-    batch_btn.click(process_zip, zip_file, 
-                    [output_acc, output_roc, output_report, output_cm, output_roc_plot])
-
-app.launch(show_api=False, max_threads=24)
+        message = f"Image not Found<br>Error: {e}"
+    return image, message
+
+detector = AIDetector()
+
+def create_gradio_interface():
+    with gr.Blocks() as app:
+        gr.Markdown("""<center><h1>AI Image Detector<br><h4>(Test Demo - accuracy varies by model)</h4></h1></center>""")
+
+        with gr.Tabs():
+            with gr.Tab("Single Image Detection"):
+                with gr.Column():
+                    inp = gr.Image(type='pil')
+                    in_url = gr.Textbox(label="Image URL")
+                    with gr.Row():
+                        load_btn = gr.Button("Load URL")
+                        btn = gr.Button("Detect AI")
+                    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.Tab("Batch Image Processing"):
+                zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
+                batch_btn = gr.Button("Process Batch")
+
+                with gr.Group():
+                    gr.Markdown(f"### Results for {MODEL_NAME}")
+                    output_acc = gr.Label(label="Accuracy")
+                    output_roc = gr.Label(label="ROC Score")
+                    output_report = gr.Textbox(label="Classification Report", lines=10)
+                    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(
+            process_zip,
+            zip_file,
+            [output_acc, output_roc, output_report, output_plots]
+        )
+
+    return app
+
+if __name__ == "__main__":
+    app = create_gradio_interface()
+    app.launch(show_api=False, max_threads=24)