Update vit_model_test.py

vit_model_test.py

@@ -1,15 +1,32 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 from torch.utils.data import Dataset, DataLoader
 from torchvision import transforms
 from transformers import ViTForImageClassification
 from PIL import Image
 import os
 import pandas as pd
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, precision_score, confusion_matrix, f1_score, average_precision_score
+import matplotlib.pyplot as plt
+import seaborn as sns
+from sklearn.metrics import recall_score
 
+def labeling(path_real, path_fake):
+    image_paths = []
+    labels = []
+
+    for filename in os.listdir(path_real):
+        image_paths.append(os.path.join(path_real, filename))
+        labels.append(0)
+
+    for filename in os.listdir(path_fake):
+        image_paths.append(os.path.join(path_fake, filename))
+        labels.append(1)
 
+    dataset = pd.DataFrame({'image_path': image_paths, 'label': labels})
 
+    return dataset
 
 class CustomDataset(Dataset):
     def __init__(self, dataframe, transform=None):
@@ -26,9 +43,19 @@ class CustomDataset(Dataset):
         if self.transform:
             image = self.transform(image)
 
-        return image
+        label = self.dataframe.iloc[idx, 1]  # Label is in the second column
+        return image, label
     
-      
+def shuffle_and_split_data(dataframe, test_size=0.2, random_state=59):
+    # Shuffle the DataFrame
+    shuffled_df = dataframe.sample(frac=1, random_state=random_state).reset_index(drop=True)
+    
+    # Split the DataFrame into train and validation sets
+    train_df, val_df = train_test_split(shuffled_df, test_size=test_size, random_state=random_state)
+    
+    return train_df, val_df    
+
+
 if __name__ == "__main__":
     # Check for GPU availability
     device = torch.device('cuda')
@@ -45,25 +72,11 @@ if __name__ == "__main__":
         transforms.ToTensor()
     ])
 
-
-
-    
-
     # Load the test dataset
-    
-    
-    ### need to recive image from gratio/streamlit
-    
-    test_set = 'datasets/'
-    
-    image_paths = []
-    for filename in os.listdir(test_set):
-        image_paths.append(os.path.join(test_set, filename))
-        dataset = pd.DataFrame({'image_path': image_paths})
-        
-    
-    
-    test_dataset = CustomDataset(dataset, transform=preprocess)
+    test_real_folder = 'test/art/real'
+    test_fake_folder = 'test/art/fake'
+    test_set = labeling(test_real_folder, test_fake_folder)
+    test_dataset = CustomDataset(test_set, transform=preprocess)
     test_loader = DataLoader(test_dataset, batch_size=32)
 
     # Load the trained model
@@ -71,25 +84,37 @@ if __name__ == "__main__":
 
     # Evaluate the model
     model.eval()
-    confidences = []
+    true_labels = []
     predicted_labels = []
-    
 
     with torch.no_grad():
-        for images in test_loader:
-            images = images.to(device)
+        for images, labels in test_loader:
+            images, labels = images.to(device), labels.to(device)
             outputs = model(images)
             logits = outputs.logits  # Extract logits from the output
-            probabilities = F.softmax(logits, dim=1)
-            confidences_per_image, predicted = torch.max(probabilities, 1)
+            _, predicted = torch.max(logits, 1)
+            true_labels.extend(labels.cpu().numpy())
             predicted_labels.extend(predicted.cpu().numpy())
-            confidences.extend(confidences_per_image.cpu().numpy())
-
-
-    print(predicted_labels)
-    print(confidences)
-    
-    confidence_percentages = [confidence * 100 for confidence in confidences]
-for label, confidence in zip(predicted_labels, confidence_percentages):
-   print(f"Predicted label: {label}, Confidence: {confidence:.2f}%")
 
+    # Calculate evaluation metrics
+    accuracy = accuracy_score(true_labels, predicted_labels)
+    precision = precision_score(true_labels, predicted_labels)
+    cm = confusion_matrix(true_labels, predicted_labels)
+    f1 = f1_score(true_labels, predicted_labels)
+    ap = average_precision_score(true_labels, predicted_labels)
+    recall = recall_score(true_labels, predicted_labels)
+
+
+    print(f"Test Accuracy: {accuracy:.2%}")
+    print(f"Precision: {precision:.2%}")
+    print(f"F1 Score: {f1:.2%}")
+    print(f"Average Precision: {ap:.2%}")
+    print(f"Recall: {recall:.2%}")
+
+    # Plot the confusion matrix
+    plt.figure(figsize=(8, 6))
+    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
+    plt.xlabel('Predicted Labels')
+    plt.ylabel('True Labels')
+    plt.title('Confusion Matrix')
+    plt.show()