from ultralytics import YOLO
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["CUDA_VISIBLE_DEVICES"]="0"


if __name__ == '__main__':
    model = YOLO('yolo11n.pt')

    results = model.train(
        data = "D:\AIM\lemon\AIM.v3i.yolov11\data.yaml",       # Path to the dataset (YAML file for COCO128)
        epochs = 50,                              # Number of epochs for fine-tuning
        imgsz = 640,                             # Image size for training
        batch = 16,                              # Batch size
        lr0 = 0.005,                              # Initial learning rate
        workers = 4,                             # Number of data loading workers
        device = 0                               # Specify the device (0 = first GPU, 'cpu' for CPU)
    )