Update Vit_Traning.py

Vit_Traning.py

@@ -9,21 +9,6 @@ import os
 import pandas as pd
 from sklearn.model_selection import train_test_split
 
-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):
@@ -48,7 +33,7 @@ def shuffle_and_split_data(dataframe, test_size=0.2, random_state=59):
     train_df, val_df = train_test_split(shuffled_df, test_size=test_size, random_state=random_state)
     return train_df, val_df
 
-class CustomModel:
+class Custom_VIT_Model:
     def __init__(self):
         # Check for GPU availability
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
@@ -139,7 +124,7 @@ class CustomModel:
 
 if __name__ == "__main__":
     # Initialize the model
-    custom_model = CustomModel()
+    custom_model = Custom_VIT_Model()
 
     # Example usage: adding a new image and label
     # custom_model.add_data('path/to/image.jpg', 0)  # 0 for real, 1 for fake