Update app.py

app.py

@@ -178,47 +178,45 @@ def classify_images(image_dir, model_pipeline, model_idx):
     return labels, preds, images
 
 
-# Function to generate evaluation metrics
-def evaluate_model(labels, preds):
-    cm = confusion_matrix(labels, preds)
-    accuracy = accuracy_score(labels, preds)
-    roc_score = roc_auc_score(labels, preds)
-    report = classification_report(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
+# Function to classify images in a folder
+def classify_images(image_dir, model_pipeline, model_idx):
+    images = []
+    labels = []
+    preds = []
+    for folder_name, ground_truth_label in [('real', 1), ('ai', 0)]:
+        folder_path = os.path.join(image_dir, folder_name)
+        if not os.path.exists(folder_path):
+            continue
+        for img_name in os.listdir(folder_path):
+            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 = np.argmax([x['score'] for x in pred])
+
+                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}")
+    return labels, preds, images
 
-# Batch processing for all models
-# Batch processing for all models
 # Batch processing for all models
 def process_zip(zip_file):
     extracted_dir = extract_zip(zip_file.name)
 
-    # Initialize model pipelines (already initialized outside)
+    # Initialize model pipelines separately to avoid reuse issues
     model_pipelines = [pipe0, pipe1, pipe2]
 
     # Run classification for each model
     results = {}
-    for idx, pipe in enumerate(model_pipelines):  # Ensure each model pipeline is used separately
+    for idx, pipe in enumerate(model_pipelines):
         print(f"Processing with model {idx}")
-        labels, preds, images = classify_images(extracted_dir, pipe, idx)  # Pass in model index for debugging
+        
+        # Classify images with the correct pipeline per model
+        labels, preds, images = classify_images(extracted_dir, pipe, idx)
         accuracy, roc_score, report, cm_fig, roc_fig = evaluate_model(labels, preds)
 
         # Store results for each model
@@ -240,6 +238,7 @@ def process_zip(zip_file):
 
 
 
+
 # Single image section
 def load_url(url):
     try: