from ultralytics import YOLO
import os

# Define path to the config.yaml in your Hugging Face Space
ROOT_DIR = '/app/data'  # Since data is uploaded under /app/data in Hugging Face Space
config_path = os.path.join(ROOT_DIR, '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)