Upload train_model.py

train_model.py

@@ -0,0 +1,74 @@
+import os
+from ultralytics import YOLO
+
+# --- Configuration ---
+# Path to your data.yaml file.
+DATA_YAML_PATH = "C:\\Users\\Veera\\Downloads\\My First Project.v1i.yolov8\\data.yaml"
+
+# Choose a pre-trained YOLOv8 model to start with.
+PRETRAINED_MODEL = "yolov8n.pt"
+
+# Training parameters
+EPOCHS = 50 # Number of training epochs. Adjust based on your dataset size and desired accuracy.
+IMG_SIZE = 640 # Image size for training (as per Roboflow preprocessing).
+BATCH_SIZE = 8 # Reduced batch size since you don't have a GPU. You might need to lower it further if you encounter memory issues.
+PROJECT_NAME = "Model_trained" # Name of the project directory where results will be saved
+RUN_NAME = "best.pt" # Name of the specific run within the project directory
+
+# --- Main Training Logic ---
+def train_yolov8_model():
+    """
+    Trains a YOLOv8 model using the specified dataset and parameters.
+    The trained model (best.pt) will be saved in runs/detect/{RUN_NAME}/weights/.
+    """
+    print(f"Starting YOLOv8 model training with {PRETRAINED_MODEL}...")
+
+    # 1. Load a pre-trained YOLOv8 model
+    try:
+        model = YOLO(PRETRAINED_MODEL)
+        print(f"Successfully loaded pre-trained model: {PRETRAINED_MODEL}")
+    except Exception as e:
+        print(f"Error loading pre-trained model: {e}")
+        print("Please ensure you have an active internet connection if downloading for the first time.")
+        return
+
+    # 2. Check if data.yaml exists
+    if not os.path.exists(DATA_YAML_PATH):
+        print(f"Error: data.yaml not found at '{DATA_YAML_PATH}'.")
+        print("Please ensure the 'data.yaml' file is in the correct location.")
+        return
+
+    # 3. Train the model
+    print(f"Training model on dataset defined in: {DATA_YAML_PATH}")
+    print(f"Training for {EPOCHS} epochs with image size {IMG_SIZE} and batch size {BATCH_SIZE} on CPU...")
+    print("Training on CPU will be significantly slower.")
+
+    try:
+        results = model.train(
+            data=DATA_YAML_PATH,
+            epochs=EPOCHS,
+            imgsz=IMG_SIZE,
+            batch=BATCH_SIZE,
+            project=PROJECT_NAME,
+            name=RUN_NAME
+        )
+        print("\nTraining completed successfully!")
+
+        # The best.pt file is typically saved in runs/detect/{RUN_NAME}/weights/best.pt
+        output_weights_dir = os.path.join("runs", "detect", RUN_NAME, "weights")
+        best_pt_path = os.path.join(output_weights_dir, "best.pt")
+
+        if os.path.exists(best_pt_path):
+            print(f"Your trained model (best.pt) is saved at: {os.path.abspath(best_pt_path)}")
+            print("You can now use this .pt file for local inference or upload it to Hugging Face.")
+        else:
+            print("Warning: 'best.pt' file not found in the expected location after training.")
+            print(f"Please check the output directory: {os.path.abspath(output_weights_dir)}")
+
+    except Exception as e:
+        print(f"An error occurred during training: {e}")
+        print("Common issues: insufficient CPU memory (try reducing batch_size), incorrect data.yaml paths.")
+
+# Run the training function
+if __name__ == "__main__":
+    train_yolov8_model()