Update vit_Training.py

vit_Training.py

@@ -64,25 +64,24 @@ class Custom_VIT_Model:
         else:
             self.df = pd.DataFrame(columns=['image_path', 'label'])
 
+    def add_data(self, image_path: str, label: int):
+        # Create a new DataFrame entry
+        new_entry = pd.DataFrame({'image_path': [image_path], 'label': [label]})
+        
+        # Append the new entry to the existing DataFrame
+        self.df = pd.concat([self.df, new_entry], ignore_index=True)
+        
+        # Save the updated DataFrame to the specified CSV file
+        self.df.to_csv(self.data_file, index=False)
+
+        # Print the current state of the training data for debugging
+        print("Current training data:")
+        print(self.df)
 
-def add_data(self, image_path: str, label: int):
-    # Create a new DataFrame entry
-    new_entry = pd.DataFrame({'image_path': [image_path], 'label': [label]})
-    
-    # Append the new entry to the existing DataFrame
-    self.df = pd.concat([self.df, new_entry], ignore_index=True)
-    
-    # Save the updated DataFrame to the specified CSV file
-    self.df.to_csv(self.data_file, index=False)
-
-    # Print the current state of the training data for debugging
-    print("Current training data:")
-    print(self.df)
-
-    # Check if we have 100 images for retraining
-    if len(self.df) >= 100:
-        print("Retraining the model as we have enough data.")
-        self.retrain_model()
+        # Check if we have 100 images for retraining
+        if len(self.df) >= 100:
+            print("Retraining the model as we have enough data.")
+            self.retrain_model()