updates

script.py

@@ -1,14 +1,15 @@
 """
-This script is used to train the model for the project.
+Note: You don't need to modify this file as this script is used to train the model for the project.
+
+All of your work should be done in the data_curation.py script.
 
 You should import your main functions from the data_curation.py script and use them to prepare the dataset for training.
 
-The approved model is `yolov8m` from Ulytralytics. 
+The approved model is `yolov10m` from Ulytralytics. 
 
 Your predictions must be in a label_field called "predictions" in the dataset.
 
 See here for more details about hyperparameters for this model: https://docs.ultralytics.com/modes/train/#train-settings
-
 """
 import os
 from datetime import datetime
@@ -67,24 +68,35 @@ def train_model(training_dataset, training_config):
     """
     Train the YOLO model on the given dataset using the provided configuration.
     """
+    print("Starting the training process...")
+
+    print("Splitting the dataset...")
     four.random_split(training_dataset, {"train": training_config['train_split'], "val": training_config['val_split']})
+    print("Dataset split completed.")
 
+    print("Exporting dataset to YOLO format...")
     export_to_yolo_format(
         samples=training_dataset,
         classes=training_dataset.default_classes,
     )
+    print("Dataset export completed.")
 
+    print("Initializing the YOLO model...")
     model = YOLO("yolov10m.pt")
+    print("Model initialized.")
 
+    print("Starting model training...")
     results = model.train(
         data="dataset.yaml",
         **training_config['train_params']
     )
-    
+    print("Model training completed.")
+
     best_model_path = str(results.save_dir / "weights/best.pt")
+    print(f"Best model path: {best_model_path}")
     best_model = YOLO(best_model_path)
+    print("Best model loaded.")
 
-    return best_model
-
+    print(f"Best model saved to: {best_model_path}")
 if __name__=="__main__":
     run()