Spaces:
Sleeping
Sleeping
| import os | |
| from ultralytics import YOLO | |
| PROJECT_DIR = os.path.dirname(os.path.dirname(__file__)) | |
| # load a pretrained model (recommended for training) | |
| model = YOLO('yolov8n.pt') | |
| os.environ['KMP_DUPLICATE_LIB_OK'] = 'True' | |
| out_path = os.path.join(PROJECT_DIR, 'output') | |
| data_path = os.path.join(PROJECT_DIR, 'data', 'custom_data.yaml') | |
| # train on the pretrained model | |
| results = model.train( | |
| data=data_path, | |
| imgsz=640, | |
| epochs=2, | |
| batch=32, | |
| project=out_path, | |
| name='yolov8n_custom', | |
| save_period=2 | |
| ) | |
| # evaluate model performance on the validation set | |
| results = model.val(project=out_path) |