from ultralytics import YOLO
import torch



if __name__ == '__main__':

    # # 单纯安全帽数据集训练
    # model = YOLO('yolov8n.yaml')
    #
    # results = model.train(data='./5000_HeadHelmet.yaml', epochs=100, batch=1, patience=20, imgsz=640, save=True,
    #                       workers=0, device=0, project='5000_HeadHelmet', name='train', verbose=True)

    # model = YOLO('yolov8n_hhfs.yaml')

    # results = model.train(data='newdataset.yaml', epochs=100, batch=4, patience=20, imgsz=640, save=True,
    #                       workers=0, device=0, project='10000_newdataset', name='train', verbose=True)

    # print(torch.cuda.is_available())

    model = YOLO('./yolov8n.yaml')
    
    results = model.train(data='/home/admin-root/haotian/python哈汽锻8安全帽识别/dataset_20250428_BAC.yaml', epochs=200, batch=32, patience=20, imgsz=640, save=True,
                          workers=0, device=0, project='HelmetHeadBAC', name='train', verbose=True)

# print(model)