import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO

if __name__ == '__main__':
    model = YOLO('ultralytics/cfg/models/11/yolo11.yaml')
    #model.load('yolo11n.pt') # loading pretrain weights
    model.train(data='data/InfraRedSmallTarget.yaml',
                cache=False,
                imgsz=640,
                epochs=200,
                batch=16,
                close_mosaic=10,
                device='0',
                optimizer='SGD', # using SGD
                project='runs/train-obb',
                name='exp',
                )

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO


if __name__ == '__main__':
    model = YOLO('runs/train/exp/weights/best.pt')
    model.val(data='data/InfraRedSmallTarget.yaml',
              split='val',
              imgsz=640,
              batch=16,
              iou=0.6,
              rect=False,
              save_json=False,
              project='runs/val',
              name='exp',
              )

import warnings
warnings.filterwarnings('ignore')
from ultralytics import YOLO


if __name__ == '__main__':
    model = YOLO('runs/train/exp/weights/best.pt')
    model.predict(source='data/images/',
                  imgsz=640,
                  project='runs/detect',
                  name='exp',
                  save=True,
                  conf=0.2,
                  iou=0.7,
                )