import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

from ultralytics import YOLO
from ultralytics import settings


if __name__ == '__main__':
    settings.update({'wandb': False})
    # model_path = 'D:\\codes\\diffusers\\runs\\obb\\train5\\weights\\best.pt'
    # model_path = 'D:\\train_cad_detection\\runs\\obb\\train7\\weights\\best.pt'
    model = YOLO('yolov8n-obb-new.yaml')#.load(model_path)
    results = model.train(data='layouts.yaml', epochs=100, imgsz=512, batch=32, mosaic=0, translate=0, scale=0,
                          erasing=0, degrees=0, flipud=0.5, fliplr=0.5, device=[7], optimizer='Adam', pretrained=False)#, hsv_h=0, hsv_s=0, hsv_v=0
    model.export(format='onnx')
