YOLO - Detection 실습

1. 데이터셋 준비

• roboflow의 'Aquarium' 데이터셋 다운로드

mkdir aquarium
cd aquarium

curl -L "다운로드URL" > roboflow.zip
unzip roboflow.zip
rm roboflow.zip

cd ..

• data.yml 파일 내 경로 정보 수정

path: ../aquarium/
train: train/images
val: valid/images
test: test/images
(후략)

2. 훈련

• train() 함수에 경로 정보를 수정한 data.yml 파일을 입력시켜 훈련 실시

from ultralytics import YOLO

model = YOLO("yolov8n.pt")
model.train(data="aquarium/data.yaml")
      
# Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
# 1/100      2.49G      1.569      3.807      1.296        179        640: 100%|██████████| 28/28 [00:12<00:00,  2.24it/s]
# Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 4/4 [00:02<00:00,  1.43it/s]                   all        127        909     0.0143      0.454     0.0726     0.0383

# Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
# 2/100      2.58G      1.593      2.601      1.265        243        640: 100%|██████████| 28/28 [00:08<00:00,  3.20it/s]
# Class     Images  Instances      Box(P          R      mAP50  mAP50-95): 100%|██████████| 4/4 [00:02<00:00,  1.36it/s]                   all        127        909      0.328     0.0776       0.22      0.118
# (후략)

• runs/detection/train 폴더 내 results.png 파일을 통한 훈련 경과 확인
  

3. 예측

• 가장 좋은 성능을 보였던 모델 활용, 샘플 이미지로 Detection 수행

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

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

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