Update app.py

app.py

@@ -6,7 +6,6 @@ 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 tqdm import tqdm
 from PIL import Image
 import uuid
 import tempfile
@@ -16,16 +15,9 @@ import numpy as np
 from sklearn.metrics import ConfusionMatrixDisplay
 import urllib.request
 
-# Define models
-models = [
-    "umm-maybe/AI-image-detector",
-    "Organika/sdxl-detector",
-    "cmckinle/sdxl-flux-detector",
-]
-
-pipe0 = pipeline("image-classification", f"{models[0]}")
-pipe1 = pipeline("image-classification", f"{models[1]}")
-pipe2 = pipeline("image-classification", f"{models[2]}")
+# Define model
+model = "cmckinle/sdxl-flux-detector"
+pipe = pipeline("image-classification", model)
 
 fin_sum = []
 uid = uuid.uuid4()
@@ -35,98 +27,25 @@ def softmax(vector):
     e = exp(vector - vector.max())  # for numerical stability
     return e / e.sum()
 
-# Single image classification functions
-def image_classifier0(image):
+# Single image classification function
+def image_classifier(image):
     labels = ["AI", "Real"]
-    outputs = pipe0(image)
+    outputs = pipe(image)
     results = {}
     for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
+        results[labels[idx]] = float(outputs[idx]['score'])
     fin_sum.append(results)
     return results
 
-def image_classifier1(image):
-    labels = ["AI", "Real"]
-    outputs = pipe1(image)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
-    fin_sum.append(results)
-    return results
-
-def image_classifier2(image):
-    labels = ["AI", "Real"]
-    outputs = pipe2(image)
-    results = {}
-    for idx, result in enumerate(outputs):
-        results[labels[idx]] = float(outputs[idx]['score'])  # Convert to float
-    fin_sum.append(results)
-    return results
-
-def aiornot0(image):
-    labels = ["AI", "Real"]
-    mod = models[0]
-    feature_extractor0 = AutoFeatureExtractor.from_pretrained(mod)
-    model0 = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor0(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model0(**input)
-        logits = outputs.logits
-        probability = softmax(logits)  # Apply softmax on 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
-
-def aiornot1(image):
+def aiornot(image):
     labels = ["AI", "Real"]
-    mod = models[1]
-    feature_extractor1 = AutoFeatureExtractor.from_pretrained(mod)
-    model1 = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor1(image, return_tensors="pt")
+    feature_extractor = AutoFeatureExtractor.from_pretrained(model)
+    model_cls = AutoModelForImageClassification.from_pretrained(model)
+    input = feature_extractor(image, return_tensors="pt")
     with torch.no_grad():
-        outputs = model1(**input)
+        outputs = model_cls(**input)
         logits = outputs.logits
-        probability = softmax(logits)  # Apply softmax on 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
-
-def aiornot2(image):
-    labels = ["AI", "Real"]
-    mod = models[2]
-    feature_extractor2 = AutoFeatureExtractor.from_pretrained(mod)
-    model2 = AutoModelForImageClassification.from_pretrained(mod)
-    input = feature_extractor2(image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model2(**input)
-        logits = outputs.logits
-        probability = softmax(logits)  # Apply softmax on logits
+        probability = softmax(logits)
         px = pd.DataFrame(probability.numpy())
     prediction = logits.argmax(-1).item()
     label = labels[prediction]
@@ -146,13 +65,13 @@ def aiornot2(image):
 
 # Function to extract images from zip
 def extract_zip(zip_file):
-    temp_dir = tempfile.mkdtemp()  # Temporary directory
+    temp_dir = tempfile.mkdtemp()
     with zipfile.ZipFile(zip_file, 'r') as z:
         z.extractall(temp_dir)
     return temp_dir
 
 # Function to classify images in a folder
-def classify_images(image_dir, model_pipeline, model_idx):
+def classify_images(image_dir):
     images = []
     labels = []
     preds = []
@@ -165,18 +84,16 @@ def classify_images(image_dir, model_pipeline, model_idx):
             img_path = os.path.join(folder_path, img_name)
             try:
                 img = Image.open(img_path).convert("RGB")
-                
-                # Ensure that each image is being processed by the correct model pipeline
-                pred = model_pipeline(img)
-                pred_label = 0 if pred[0]['label'] == 'AI' else 1  # Assuming 'AI' is label 0 and 'Real' is label 1
+                pred = pipe(img)
+                pred_label = 0 if pred[0]['label'] == 'AI' 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} in model {model_idx}: {e}")
+                print(f"Error processing image {img_name}: {e}")
     
-    print(f"Model {model_idx} processed {len(images)} images")
+    print(f"Processed {len(images)} images")
     return labels, preds, images
 
 # Function to generate evaluation metrics
@@ -206,46 +123,13 @@ def evaluate_model(labels, preds):
 
     return accuracy, roc_score, report, fig, fig_roc
 
-# Batch processing for all models
+# Batch processing
 def process_zip(zip_file):
     extracted_dir = extract_zip(zip_file.name)
-
-    # Run classification for each model
-    results = {}
-    for idx in range(len(models)):
-        print(f"Processing with model {models[idx]}")  # Debugging to show which model is being used
-
-        # Create a new pipeline for each model within the loop
-        pipe = pipeline("image-classification", f"{models[idx]}")  
-        print(f"Initialized pipeline for {models[idx]}")  # Confirm pipeline is initialized correctly
-
-        # Classify images with the correct pipeline per model
-        labels, preds, images = classify_images(extracted_dir, pipe, idx)
-        
-        # Debugging: Print the predictions to ensure they're different
-        print(f"Predictions for model {models[idx]}: {preds}")
-
-        accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds)
-
-        # Store results for each model
-        results[f'Model_{idx}_accuracy'] = accuracy
-        results[f'Model_{idx}_roc_score'] = roc_score
-        results[f'Model_{idx}_report'] = report
-        results[f'Model_{idx}_cm_fig'] = cm_fig
-        results[f'Model_{idx}_roc_fig'] = roc_fig
-
+    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 results for all models
-    return (results['Model_0_accuracy'], results['Model_0_roc_score'], results['Model_0_report'], 
-            results['Model_0_cm_fig'], results['Model_0_roc_fig'],
-           results['Model_1_accuracy'], results['Model_1_roc_score'], results['Model_1_report'], 
-            results['Model_1_cm_fig'], results['Model_1_roc_fig'],
-           results['Model_2_accuracy'], results['Model_2_roc_score'], results['Model_2_report'], 
-            results['Model_2_cm_fig'], results['Model_2_roc_fig'])
-
-
-
+    return accuracy, roc_score, report, cm_fig, roc_fig
 
 # Single image section
 def load_url(url):
@@ -294,18 +178,15 @@ with gr.Blocks() as app:
                 with gr.Row():
                     fin = gr.Label(label="Final Probability")
                 with gr.Row():
-                    for i, model in enumerate(models):
-                        with gr.Box():
-                            gr.HTML(f"""<b>Testing on Model {i}: <a href='https://huggingface.co/{model}'>{model}</a></b>""")
-                            globals()[f'outp{i}'] = gr.HTML("""""")
-                            globals()[f'n_out{i}'] = gr.Label(label="Output")
+                    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(aiornot0, [inp], [outp0, n_out0]).then(
-                aiornot1, [inp], [outp1, n_out1]).then(
-                aiornot2, [inp], [outp2, n_out2]).then(
+            btn.click(aiornot, [inp], [outp, n_out]).then(
                 tot_prob, None, fin, show_progress=False)
 
         # Tab for batch processing
@@ -313,19 +194,16 @@ with gr.Blocks() as app:
             zip_file = gr.File(label="Upload Zip (two folders: real, ai)")
             batch_btn = gr.Button("Process Batch")
 
-            for i, model in enumerate(models):
-                with gr.Group():
-                    gr.Markdown(f"### Results for {model}")
-                    globals()[f'output_acc{i}'] = gr.Label(label=f"Model {i} Accuracy")
-                    globals()[f'output_roc{i}'] = gr.Label(label=f"Model {i} ROC Score")
-                    globals()[f'output_report{i}'] = gr.Textbox(label=f"Model {i} Classification Report", lines=10)
-                    globals()[f'output_cm{i}'] = gr.Plot(label=f"Model {i} Confusion Matrix")
-                    globals()[f'output_roc_plot{i}'] = gr.Plot(label=f"Model {i} ROC Curve")
+            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_acc0, output_roc0, output_report0, output_cm0, output_roc_plot0,
-                     output_acc1, output_roc1, output_report1, output_cm1, output_roc_plot1,
-                     output_acc2, output_roc2, output_report2, output_cm2, output_roc_plot2])
+                    [output_acc, output_roc, output_report, output_cm, output_roc_plot])
 
-app.launch(show_api=False, max_threads=24)
+app.launch(show_api=False, max_threads=24)