YOLO - Segmentation 실습

1. 훈련

• YOLO 제공 coco8-seg.yaml 활용 Segmentation 모델 훈련

from ultralytics import YOLO

model = YOLO("yolov8n-seg.pt")
res = model.train(
	data="coco8-seg.yaml", # 데이터 부재 시 자동 다운로드
    epochs=50,
    imgsz=640
)

• runs/segment/train 폴더 내 results.png 를 통해 훈련 결과 확인
  

2. 예측

• 가장 성능이 좋았던 모델을 기반으로, 샘플 이미지에 대한 예측 수행

best_model = YOLO("runs/segment/train/weights/best.pt")
bets_model.predict("sample.jpg", save=True)

import matplotlib.pyplot as plt
img = plt.imread("runs/segment/predict/sample.jpg")
plt.imshow(img)
plt.axis("off")
plt.show()

*이 글은 제로베이스 데이터 취업 스쿨의 강의 자료 일부를 발췌하여 작성되었습니다.