Added functionalities

pages/Model_Evaluation.py

@@ -1,59 +1,78 @@
-import streamlit as st
+import sys
+import os
+import time
+import cv2
+import numpy as np
+import pandas as pd
+import seaborn as sns
 import torch
+import torch.nn as nn
 from torch.utils.data import DataLoader, Dataset
 from torchvision import transforms, models
-import torch.nn as nn
 from PIL import Image
-import pandas as pd
-import os
-import cv2
-import numpy as np
+from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
+from sklearn.preprocessing import label_binarize
+import streamlit as st
+import matplotlib.pyplot as plt
+from fpdf import FPDF
 
+# ---- Stop Evaluation button -----
+if 'stop_eval' not in st.session_state:
+    st.session_state.stop_eval = False
+
+if 'evaluation_done' not in st.session_state:
+    st.session_state.evaluation_done = False
+
+# ---- Streamlit Title ----
 st.markdown("<h2 style='color: #2E86C1;'>📈 Model Evaluation</h2>", unsafe_allow_html=True)
 
-# Define class names and label map
+# ---- Class Names & Label Mapping ----
 class_names = ['No_DR', 'Mild', 'Moderate', 'Severe', 'Proliferative_DR']
 label_map = {label: idx for idx, label in enumerate(class_names)}
 
-# Define your image preprocessing functions
+# ---- Text Cleaning Function for PDF ----
+def clean_text(text):
+    return text.encode('utf-8', 'ignore').decode('utf-8')
+
+# ---- Preprocessing Functions ----
 def apply_median_filter(image):
     return cv2.medianBlur(image, 5)
 
 def apply_clahe(image):
     lab = cv2.cvtColor(image, cv2.COLOR_RGB2LAB)
     l, a, b = cv2.split(lab)
-    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    clahe = cv2.createCLAHE(clipLimit=2.0)
     cl = clahe.apply(l)
     merged = cv2.merge((cl, a, b))
     return cv2.cvtColor(merged, cv2.COLOR_LAB2RGB)
 
-def apply_gamma_correction(image, gamma=1.5):
+def apply_gamma_correction(image, gamma=1.2):
     invGamma = 1.0 / gamma
-    table = np.array([(i / 255.0) ** invGamma * 255 for i in np.arange(256)]).astype("uint8")
+    table = np.array([(i / 255.0) ** invGamma * 255 for i in np.arange(0, 256)]).astype("uint8")
     return cv2.LUT(image, table)
 
 def apply_gaussian_filter(image, kernel_size=(5, 5), sigma=1.0):
     return cv2.GaussianBlur(image, kernel_size, sigma)
 
-# Custom dataset with preprocessing
+# ---- Custom Dataset ----
 class DDRDataset(Dataset):
-    def __init__(self, csv_path, img_dir, transform=None):
+    def __init__(self, csv_path, transform=None):
         self.data = pd.read_csv(csv_path)
-        self.img_dir = img_dir
+        self.image_paths = self.data['new_path'].tolist()
+        self.labels = self.data['label'].tolist()
         self.transform = transform
 
     def __len__(self):
-        return len(self.data)
+        return len(self.image_paths)
 
     def __getitem__(self, idx):
-        img_name = self.data.iloc[idx, 0]
-        label_name = self.data.iloc[idx, 1]
-        label = int(label_map.get(label_name, 0))  # fallback to 0
+        img_path = self.image_paths[idx]
+        label = int(self.labels[idx])
 
-        img_path = os.path.join(self.img_dir, img_name)
         image = cv2.imread(img_path)
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
+        # Apply preprocessing
         image = apply_median_filter(image)
         image = apply_clahe(image)
         image = apply_gamma_correction(image)
@@ -63,67 +82,166 @@ class DDRDataset(Dataset):
         if self.transform:
             image = self.transform(image)
 
-        return image, label
+        return image, torch.tensor(label, dtype=torch.long)
 
-# -------------------------------
-# Load Test Data with Caching
-# -------------------------------
+# ---- Image Transforms ----
+val_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# ---- Load Data (with caching) ----
 @st.cache_resource
-def load_test_data():
-    transform = transforms.Compose([
-        transforms.Resize((224, 224)),
-        transforms.ToTensor(),
-        transforms.Normalize([0.485, 0.456, 0.406],
-                             [0.229, 0.224, 0.225])
-    ])
-    dataset = DDRDataset(
-        csv_path="D:/DR_Classification/splits/test_labels.csv",
-        img_dir="D:/DR_Classification/splits/test",
-        transform=transform
-    )
+def load_test_data(csv_path):
+    dataset = DDRDataset(csv_path=csv_path, transform=val_transform)
     return DataLoader(dataset, batch_size=32, shuffle=False)
 
-# -------------------------------
-# Evaluation Function
-# -------------------------------
-def evaluation_test_model(model, test_loader, criterion, device='cpu'):
+# ---- Load Model (with caching) ----
+@st.cache_resource
+def load_model():
+    model = models.densenet121(pretrained=False)
+    model.classifier = nn.Linear(model.classifier.in_features, len(class_names))
+    model.load_state_dict(torch.load(r"D:\\DR_Classification\\training\\Pretrained_Densenet-121.pth", map_location=torch.device('cpu')))
     model.eval()
-    running_loss = 0.0
-    correct = 0
-    total = 0
+    return model
 
-    with torch.no_grad():
-        for inputs, labels in test_loader:
-            inputs, labels = inputs.to(device), labels.to(device)
-            outputs = model(inputs)
-            loss = criterion(outputs, labels)
+# ---- Main Evaluation ----
+csv_path = r"D:\\DR_Classification\\splits\\test_labels.csv"
+model = load_model()
+test_loader = load_test_data(csv_path)
 
-            running_loss += loss.item()
-            _, predicted = torch.max(outputs, 1)
-            correct += (predicted == labels).sum().item()
-            total += labels.size(0)
+col1, col2 = st.columns([1, 1])
+
+with col1:
+    if st.button("🚀 Start Evaluation"):
+        st.session_state.stop_eval = False
+        st.session_state.evaluation_done = False
+        run_eval = True
+    else:
+        run_eval = False
 
-    val_loss = running_loss / len(test_loader)
-    val_acc = correct / total * 100
-    return val_loss, val_acc
+with col2:
+    if st.button("🚩 Stop Evaluation"):
+        st.session_state.stop_eval = True
 
-# -------------------------------
-# Evaluation Button Trigger
-# -------------------------------
-if st.button("🔍 Evaluate Trained Model"):
-    with st.spinner("Evaluating on test data..."):
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+if st.session_state.evaluation_done:
+    reevaluate_col, download_col = st.columns([1, 1])
 
-        test_loader = load_test_data()
+if run_eval or st.session_state.evaluation_done:
+    st.markdown("### ⏱️ Evaluation Results")
 
-        model = models.densenet121(pretrained=False)
-        model.classifier = nn.Linear(model.classifier.in_features, len(class_names))
-        model.load_state_dict(torch.load("training/Pretrained_Densenet-121.pth", map_location=device))
-        model.to(device)
+    start_time = time.time()
+    y_true = []
+    y_pred = []
+    y_score = []
+    misclassified_images = []
 
-        criterion = nn.CrossEntropyLoss()
-        val_loss, val_acc = evaluation_test_model(model, test_loader, criterion, device)
+    total_batches = len(test_loader)
+    progress_bar = st.progress(0)
+    status_text = st.empty()
+    stop_info = st.empty()
 
-        st.success("✅ Evaluation Complete")
-        st.metric("Test Loss", f"{val_loss:.4f}")
-        st.metric("Test Accuracy", f"{val_acc:.2f}%")
+    with torch.no_grad():
+        for i, (images, labels) in enumerate(test_loader):
+            if st.session_state.stop_eval:
+                stop_info.warning("🚩 Evaluation stopped by user.")
+                break
+
+            outputs = model(images)
+            _, predicted = torch.max(outputs, 1)
+            y_true.extend(labels.numpy())
+            y_pred.extend(predicted.numpy())
+            y_score.extend(outputs.detach().numpy())
+
+            for j in range(len(labels)):
+                if predicted[j] != labels[j]:
+                    misclassified_images.append((images[j], predicted[j].item(), labels[j].item()))
+
+            percent_complete = (i + 1) / total_batches
+            progress_bar.progress(min(percent_complete, 1.0))
+            status_text.text(f"Evaluating on Test Set: {int(percent_complete * 100)}% | Batch {i+1}/{total_batches}")
+            time.sleep(0.1)
+
+    end_time = time.time()
+    eval_time = end_time - start_time
+
+    if not st.session_state.stop_eval:
+        st.session_state.evaluation_done = True
+        st.success(f"✅ Evaluation completed in **{eval_time:.2f} seconds**")
+
+        report = classification_report(y_true, y_pred, target_names=class_names, output_dict=True)
+        report_df = pd.DataFrame(report).transpose()
+        st.dataframe(report_df.style.format("{:.2f}"))
+
+        pdf = FPDF()
+        pdf.add_page()
+        pdf.set_font("Arial", size=12)
+        pdf.cell(200, 10, txt=clean_text("Classification Report"), ln=True, align='C')
+
+        col_widths = [40, 40, 40, 40]
+        headers = ["Class", "Precision", "Recall", "F1-Score"]
+        for i, header in enumerate(headers):
+            pdf.cell(col_widths[i], 10, header, border=1)
+        pdf.ln()
+
+        for idx, row in report_df.iterrows():
+            if idx in ['accuracy', 'macro avg', 'weighted avg']:
+                continue
+            pdf.cell(col_widths[0], 10, str(idx), border=1)
+            pdf.cell(col_widths[1], 10, f"{row['precision']:.2f}", border=1)
+            pdf.cell(col_widths[2], 10, f"{row['recall']:.2f}", border=1)
+            pdf.cell(col_widths[3], 10, f"{row['f1-score']:.2f}", border=1)
+            pdf.ln()
+
+        cm = confusion_matrix(y_true, y_pred)
+        fig_cm, ax = plt.subplots()
+        sns.heatmap(cm, annot=True, fmt='d', xticklabels=class_names, yticklabels=class_names, cmap="Blues", ax=ax)
+        ax.set_xlabel('Predicted')
+        ax.set_ylabel('True')
+        ax.set_title("Confusion Matrix")
+        st.pyplot(fig_cm)
+        cm_path = "confusion_matrix.png"
+        fig_cm.savefig(cm_path, format='png', dpi=300, bbox_inches='tight')
+        plt.close(fig_cm)
+        if os.path.exists(cm_path):
+            pdf.image(cm_path, x=10, y=None, w=180)
+
+        y_true_bin = label_binarize(y_true, classes=list(range(len(class_names))))
+        y_score_np = np.array(y_score)
+        fig_roc, ax = plt.subplots()
+        for i in range(len(class_names)):
+            fpr, tpr, _ = roc_curve(y_true_bin[:, i], y_score_np[:, i])
+            roc_auc = auc(fpr, tpr)
+            ax.plot(fpr, tpr, label=f'{class_names[i]} (AUC = {roc_auc:.2f})')
+
+        ax.plot([0, 1], [0, 1], 'k--')
+        ax.set_xlabel('False Positive Rate')
+        ax.set_ylabel('True Positive Rate')
+        ax.set_title('Multi-class ROC Curve')
+        ax.legend(loc='lower right')
+        st.pyplot(fig_roc)
+        roc_path = "roc_curve.png"
+        fig_roc.savefig(roc_path, format='png', dpi=300, bbox_inches='tight')
+        plt.close(fig_roc)
+        if os.path.exists(roc_path):
+            pdf.image(roc_path, x=10, y=None, w=180)
+
+        st.markdown("### ❌ Misclassified Samples")
+        fig_mis, axs = plt.subplots(1, min(5, len(misclassified_images)), figsize=(15, 4))
+        for idx, (img, pred, true) in enumerate(misclassified_images[:5]):
+            axs[idx].imshow(img.permute(1, 2, 0))
+            axs[idx].set_title(f"True: {class_names[true]}\nPred: {class_names[pred]}")
+            axs[idx].axis('off')
+        st.pyplot(fig_mis)
+
+        output_pdf = "evaluation_report.pdf"
+        pdf.output(output_pdf)
+        with open(output_pdf, "rb") as f:
+            reevaluate_col, download_col = st.columns([1, 1])
+            with download_col:
+                st.download_button("📄 Download Full Evaluation PDF", f, file_name="evaluation_report.pdf")
+            with reevaluate_col:
+                if st.button("🔁 Re-evaluate"):
+                    st.session_state.stop_eval = False
+                    st.session_state.evaluation_done = False