Update vit_Training.py

vit_Training.py

@@ -19,13 +19,13 @@ class CustomDataset(Dataset):
         return len(self.dataframe)
 
     def __getitem__(self, idx):
-        image_path = self.dataframe.iloc[idx, 0]  # Image path is in the first column
+        image_path = self.dataframe.iloc[idx, 0]  
         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
+        label = self.dataframe.iloc[idx, 1]  
         return image, label
 
 def shuffle_and_split_data(dataframe, test_size=0.2, random_state=59):
@@ -35,23 +35,23 @@ def shuffle_and_split_data(dataframe, test_size=0.2, random_state=59):
 
 class Custom_VIT_Model:
     def __init__(self):
-        # Check for GPU availability
+        # Use gpu if exist (nvidia only) else cpu (any)
         self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
         
-        # Load the pre-trained ViT model and move it to the device
+        # Load the pre-trained ViT model
         self.model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224').to(self.device)
 
         # Freeze pre-trained layers
         for param in self.model.parameters():
             param.requires_grad = False
 
-        # Define a new classifier and move it to the device
+        # Define a new classifier that has 2 outputs (0,1)
         self.model.classifier = nn.Linear(self.model.config.hidden_size, 2).to(self.device)
 
-        # Define the optimizer
+        # Set optimizer
         self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
 
-        # Define the image preprocessing pipeline
+        # Set the image preprocessing (resize image) and make it tensor ( Tensor - add a dimension )
         self.preprocess = transforms.Compose([
             transforms.Resize((224, 224)),
             transforms.ToTensor()
@@ -103,7 +103,7 @@ class Custom_VIT_Model:
         # Define the loss function
         criterion = nn.CrossEntropyLoss().to(self.device)
 
-        # Training loop
+        # Training loop 
         num_epochs = 10
         for epoch in range(num_epochs):
             self.model.train()
@@ -139,8 +139,5 @@ class Custom_VIT_Model:
         print("Model retrained and updated!")
 
 if __name__ == "__main__":
-    # Initialize the model
     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