Spaces:
Sleeping
Sleeping
File size: 3,146 Bytes
3c0ad87 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import os
from ultralytics import YOLO
# --- Configuration ---
# Path to your data.yaml file.
DATA_YAML_PATH = "C:\Users\Veera\Downloads\Yolov8n-env\data.yaml"
# Choose a pre-trained YOLOv8 model to start with.
PRETRAINED_MODEL = "yolov8n.pt"
# Training parameters
EPOCHS = 100 # 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 = 16 # 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() |