from ultralytics import YOLO

def train_model(model_path, data_path, epochs, batch_size, device, img_size, save_dir):
    """
    Trains the YOLO model on the specified dataset.

    Parameters:
    model_path (str): Path to the YOLO model.
    data_path (str): Path to the dataset configuration file.
    epochs (int): Number of training epochs.
    batch_size (int): Batch size for training.
    device (str): Device to use for training ('cpu' or 'cuda').
    img_size (int): Image size for training.
    save_dir (str): Directory to save the trained model.
    """
    model = YOLO('yolov8n.pt')
    model.train(
        data=data_path,
        epochs=epochs,
        batch=batch_size,
        device=device,
        imgsz=img_size,
        cache=True,
        save_dir=save_dir
    )

if __name__ == "__main__":
    train_model(
        model_path='yolov8n.pt',
        data_path='Dataset/datasets.yaml',
        epochs=100,
        batch_size=16,
        device='cuda',
        img_size=320,
        save_dir='Dataset/model'
    )