File size: 669 Bytes
bb14d6a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | import os
from ultralytics import YOLO
project = f"{os.getcwd()}/runs"
device = "cuda:2"
imgsz = 640
epochs = 20
# Stage 1: Pretrain on SeaTurtleID2022 (large dataset)
model = YOLO("yolo11s-seg.pt")
model.train(
data="segmentation_stage1.yaml",
project=project,
name="stage1",
epochs=epochs,
imgsz=imgsz,
device=device,
fliplr=0,
flipud=0,
)
# Stage 2: Fine-tune on combined dataset (balanced)
model = YOLO(f"{project}/stage1/weights/last.pt")
model.train(
data="segmentation_stage2.yaml",
project=project,
name="stage2",
epochs=epochs,
imgsz=imgsz,
device=device,
fliplr=0,
flipud=0,
freeze=5,
)
|