app.py

from ultralytics import YOLO
import os

# Define the correct path to config.yaml (in the root directory)
config_path = './config.yaml'  # Adjust based on the actual path to your config.yaml

# Load YOLO model
model = YOLO("yolo11n.yaml")  # You can choose a different model type like yolo5n, yolo6n, etc.

# Train the model
results = model.train(data=config_path, epochs=1)

# Define the save directory
save_dir = './runs/detect/train/weights'

# Create directory if it doesn't exist
if not os.path.exists(save_dir):
    os.makedirs(save_dir)

# Save the model
model.save(os.path.join(save_dir, 'best.pt'))

# Print confirmation
print("Model saved to:", os.path.join(save_dir, 'best.pt'))