update

vit_model_test.py

@@ -1,32 +1,15 @@
 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):
@@ -43,19 +26,9 @@ class CustomDataset(Dataset):
         if self.transform:
             image = self.transform(image)
 
-        label = self.dataframe.iloc[idx, 1]  # Label is in the second column
-        return image, label
+        return image
     
-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')
@@ -72,11 +45,25 @@ if __name__ == "__main__":
         transforms.ToTensor()
     ])
 
+
+
+    
+
     # Load the test dataset
-    test_real_folder = 'datasets/test_set/real'
-    test_fake_folder = 'datasets/test_set/fake'
-    test_set = labeling(test_real_folder, test_fake_folder)
-    test_dataset = CustomDataset(test_set, transform=preprocess)
+    
+    
+    ### 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_loader = DataLoader(test_dataset, batch_size=32)
 
     # Load the trained model
@@ -84,37 +71,25 @@ if __name__ == "__main__":
 
     # Evaluate the model
     model.eval()
-    true_labels = []
+    confidences = []
     predicted_labels = []
+    
 
     with torch.no_grad():
-        for images, labels in test_loader:
-            images, labels = images.to(device), labels.to(device)
+        for images in test_loader:
+            images = images.to(device)
             outputs = model(images)
             logits = outputs.logits  # Extract logits from the output
-            _, predicted = torch.max(logits, 1)
-            true_labels.extend(labels.cpu().numpy())
+            probabilities = F.softmax(logits, dim=1)
+            confidences_per_image, predicted = torch.max(probabilities, 1)
             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()