revert max count to max class

train_main.py

@@ -6,18 +6,17 @@ from src import data, features, model, augmentation
 
 
 def process_dataset_with_augmentation(df, is_training=False):
-    """
-    Loops through the dataframe.
-    If is_training=True, it augments minority classes to balance the data.
-    """
+    
+    #Loops through the dataframe.
+    #If is_training=True, it augments minority classes to balance the data.
+    
     X = []
     y = []
 
     # 1. Calculate Statistics for Balancing (Only needed for training)
     if is_training:
         class_counts = df['target'].value_counts().to_dict()
-        max_count = 500
-        print(f"  [Augmentation] Balancing classes to match majority count: {max_count}")
+        max_count = max(class_counts.values())        print(f"  [Augmentation] Balancing classes to match majority count: {max_count}")
 
     total = len(df)
 
@@ -59,9 +58,8 @@ def process_dataset_with_augmentation(df, is_training=False):
 
 
 def main():
-    """
-    This is the starting point of the training.
-    It produces -between others- plots and a model saved in serialized format (classes.pkl) """
+    #This is the starting point of the training.
+    #It produces -between others- plots and a model saved in serialized format (classes.pkl) 
     # 1. Load Data (Metadata only)
     df, classes = data.load_metadata(limit=None)  # Adjust limit as needed