from ultralytics import YOLO
import os

BASE = os.path.dirname(os.path.abspath(__file__))
YOLO_DIR = os.path.join(BASE, "datasets/yolo")
data = os.path.join(YOLO_DIR, "dataset_local.yaml")
with open(data, "w", encoding="utf-8") as f:
    f.write(f"""path: {YOLO_DIR}
train: images/train
val: images/val
names:
  0: pheasant
  1: suspect
""")

model = YOLO("yolov8n.pt")
model.train(data=data, imgsz=1024, epochs=150, patience=30, batch=8,
            device="mps", workers=4,
            # project 必须绝对路径，相对路径会被拼到默认 runs/detect 下造成双层嵌套
            project=os.path.join(BASE, "runs/pheasant"), name="yolov8n_1024",
            exist_ok=True, plots=True)