Update app.py

app.py

@@ -151,33 +151,6 @@ def extract_zip(zip_file):
         z.extractall(temp_dir)
     return temp_dir
 
-# Function to classify images in a folder
-# 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")
-
-                # Now use the specific model pipeline passed in
-                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
-
-
 # Function to classify images in a folder
 def classify_images(image_dir, model_pipeline, model_idx):
     images = []
@@ -203,6 +176,33 @@ def classify_images(image_dir, model_pipeline, model_idx):
                 print(f"Error processing image {img_name} in model {model_idx}: {e}")
     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
+
 # Batch processing for all models
 def process_zip(zip_file):
     extracted_dir = extract_zip(zip_file.name)
@@ -237,8 +237,6 @@ def process_zip(zip_file):
             results['Model_2_cm_fig'], results['Model_2_roc_fig'])
 
 
-
-
 # Single image section
 def load_url(url):
     try: