22 lines
701 B
Python
22 lines
701 B
Python
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)
|