update

vit_model_test.py

@@ -11,40 +11,9 @@ from sklearn.metrics import accuracy_score, precision_score, confusion_matrix, f
 import matplotlib.pyplot as plt
 import seaborn as sns
 from sklearn.metrics import recall_score
+from vit_model_traning import labeling,CustomDataset
 
-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):
-        self.dataframe = dataframe
-        self.transform = transform
-
-    def __len__(self):
-        return len(self.dataframe)
-
-    def __getitem__(self, idx):
-        image_path = self.dataframe.iloc[idx, 0]  # Image path is in the first column
-        image = Image.open(image_path).convert('RGB')  # Convert to RGB format
-
-        if self.transform:
-            image = self.transform(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
@@ -73,8 +42,9 @@ if __name__ == "__main__":
     ])
 
     # Load the test dataset
-    test_real_folder = 'test/art/real'
-    test_fake_folder = 'test/art/fake'
+    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)
     test_loader = DataLoader(test_dataset, batch_size=32)