datascienceharp commited on
Commit
e3c4fa8
·
1 Parent(s): 8b84a69
Files changed (1) hide show
  1. script.py +18 -6
script.py CHANGED
@@ -1,14 +1,15 @@
1
  """
2
- This script is used to train the model for the project.
 
 
3
 
4
  You should import your main functions from the data_curation.py script and use them to prepare the dataset for training.
5
 
6
- The approved model is `yolov8m` from Ulytralytics.
7
 
8
  Your predictions must be in a label_field called "predictions" in the dataset.
9
 
10
  See here for more details about hyperparameters for this model: https://docs.ultralytics.com/modes/train/#train-settings
11
-
12
  """
13
  import os
14
  from datetime import datetime
@@ -67,24 +68,35 @@ def train_model(training_dataset, training_config):
67
  """
68
  Train the YOLO model on the given dataset using the provided configuration.
69
  """
 
 
 
70
  four.random_split(training_dataset, {"train": training_config['train_split'], "val": training_config['val_split']})
 
71
 
 
72
  export_to_yolo_format(
73
  samples=training_dataset,
74
  classes=training_dataset.default_classes,
75
  )
 
76
 
 
77
  model = YOLO("yolov10m.pt")
 
78
 
 
79
  results = model.train(
80
  data="dataset.yaml",
81
  **training_config['train_params']
82
  )
83
-
 
84
  best_model_path = str(results.save_dir / "weights/best.pt")
 
85
  best_model = YOLO(best_model_path)
 
86
 
87
- return best_model
88
-
89
  if __name__=="__main__":
90
  run()
 
1
  """
2
+ Note: You don't need to modify this file as this script is used to train the model for the project.
3
+
4
+ All of your work should be done in the data_curation.py script.
5
 
6
  You should import your main functions from the data_curation.py script and use them to prepare the dataset for training.
7
 
8
+ The approved model is `yolov10m` from Ulytralytics.
9
 
10
  Your predictions must be in a label_field called "predictions" in the dataset.
11
 
12
  See here for more details about hyperparameters for this model: https://docs.ultralytics.com/modes/train/#train-settings
 
13
  """
14
  import os
15
  from datetime import datetime
 
68
  """
69
  Train the YOLO model on the given dataset using the provided configuration.
70
  """
71
+ print("Starting the training process...")
72
+
73
+ print("Splitting the dataset...")
74
  four.random_split(training_dataset, {"train": training_config['train_split'], "val": training_config['val_split']})
75
+ print("Dataset split completed.")
76
 
77
+ print("Exporting dataset to YOLO format...")
78
  export_to_yolo_format(
79
  samples=training_dataset,
80
  classes=training_dataset.default_classes,
81
  )
82
+ print("Dataset export completed.")
83
 
84
+ print("Initializing the YOLO model...")
85
  model = YOLO("yolov10m.pt")
86
+ print("Model initialized.")
87
 
88
+ print("Starting model training...")
89
  results = model.train(
90
  data="dataset.yaml",
91
  **training_config['train_params']
92
  )
93
+ print("Model training completed.")
94
+
95
  best_model_path = str(results.save_dir / "weights/best.pt")
96
+ print(f"Best model path: {best_model_path}")
97
  best_model = YOLO(best_model_path)
98
+ print("Best model loaded.")
99
 
100
+ print(f"Best model saved to: {best_model_path}")
 
101
  if __name__=="__main__":
102
  run()