[YOLOv8] YOLOv8n 맛보기 + 커스텀 데이터셋

알게 된 점 위주로 작성하는 개인적인 글

 

먼저 ultralytics 레포지토리를 클론한다.

 

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite

GitHub - ultralytics/ultralytics: NEW - YOLOv8 🚀 in PyTorch > ONNX > OpenVINO > CoreML > TFLite

 

커스텀 데이터셋을 받아오기 위해 pretrain 된 모델을 받아온다.(굳이 안해도 되긴 한다)

YOLOv8 - Ultralytics YOLO 문서

YOLOv8

 

모델은 YOLOv8(n~x) 까지 있는데, 원하는 것으로 다운받으면 된다.

parameter 가 많을수록 속도가 느리지만, 정확도가 상승한다.

 

나는 가장 가벼운 YOLOv8n 을 사용한다.

 

다운받은 YOLOv8n.pt 파일을 ultralytics 폴더 안에 넣는다.

 

커스텀 데이터셋을 만들기 위해 여러가지가 가능한데, 나는 이미지를 수동으로 라벨링하고 roboflow 사이트에서 yaml 파일을 생성하는 방법으로 해 보았다. 아마 roboflow 에서 자동으로도 지원하는 것 같다.

 

수동 라벨링 툴은 labelImg 를 사용하였다.

 

YOLO 이미지 라벨링을 위한 labelImg 사용법 (velog.io)

YOLO 이미지 라벨링을 위한 labelImg 사용법

 

잘 정리된 글을 따라 가면 된다.

 

포맷을 pascalVOC 에서 YOLO 로 바꾸는 것을 까먹지 말자.

 

나는 커스텀 데이터셋으로 사과 사진을 사용하였다.

 

apple#.jpg 사진을 라벨링을 끝내고 난 파일들인데, 이를 roboflow 에 넣으면 커스텀 데이터셋을 만들어 준다.

모델을 따로 지정하지는 않고 export dataset 으로 데이터셋만 가져와 준다.

 

형식은 YOLOv8 로 지정하고 export 해 주면 위와 같이 파일들이 있다.

 

data.yaml 을 한 번 본다.

 

train: ../train/images
val: ../valid/images
test: ../test/images

nc: 1
names: ['15']

roboflow:
  workspace: hipensan
  project: apple-finder-ar2zu
  version: 1
  license: CC BY 4.0
  url: https://universe.roboflow.com/hipensan/apple-finder-ar2zu/dataset/1

 

names 에 '15' 라고 되어 있는데, 이를 객체 이름으로 바꿔준다. 나는 apple finder 를 사용하기 때문에 Apple 라고 붙여주도록 하겠다.

 

데이터셋을 준비했으니 모델에 학습시킬 차례이다.

 

ultralytics 폴더에 main.py 를 생성한다.

[main.py]

from ultralytics import YOLO

model = YOLO('yolov8n.pt')

model.train(data="../Apple Finder/data.yaml", epochs=20)

 

model.train 의 data 경로는 yaml 파일의 경로로 해 줘야 한다. 절대 경로로 하는게 정신건강에 이롭다.

 

코드 시작전 현재 command 창의 위치를 잘 확인하자. 현재 위치에 runs/.. 파일들이 생성된다.

 

+) 데이터셋이 있는 곳에서 실행시키는 것이 편할 수도 있다.

 

 

20에포크를 지나면 다음과 같이 내가 설정한 경로/runs/detect/train# 하위에 best.pt, last.pt 가 생성되었다.

 

학습시킨 best.pt의 내부 구조를 한번 본다.

 

경로 설정을 쉽게 하기 위해 ultralytics 폴더로 best.pt 파일을 복사해 온다.

 

from ultralytics import YOLO
import torch
# model = YOLO('yolov8n.pt')

# model.train(data="C:/Users/admin/Desktop/github/Apple Finder/data.yaml", epochs=20)

model = torch.load("./best.pt")
print(model)

 

레이어 0 ~ 레이어 22까지 내부 구조가 나오는 것을 확인 가능하다.

 

train 시켰으니 예측을 시켜 보는 것도 필요하겠다.

 

사과와 비슷해 보이는 복숭아, 자두, 사과 사진을 각각 준비한다.

아, 사진의 사이즈는 640*640이다. roboflow 를 사용하면 resize 도 지원한다.

 

+) 다시 위치 바꾸기 귀찮아서 데이터셋이 있는 위치에서 모두 진행하였다. best.pt 도 옮겼다.

 

파일구성은 위 사진과 같다.

 

코드는 다음과 같다. Predict - Ultralytics YOLO 문서 에서 제공해주는 코드에서 모델과 input image 만 바꿨다.

from ultralytics import YOLO

# Load a model
model = YOLO("best.pt")  # pretrained YOLOv8n model

# Run batched inference on a list of images
results = model(["peaches1.jpg", "peaches2.jpg", "peaches3.jpg", "plum1.jpg", "plum2.jpg"])  # return a list of Results objects

# Process results list
for result in results:
    boxes = result.boxes  # Boxes object for bounding box outputs
    masks = result.masks  # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs  # Probs object for classification outputs
    obb = result.obb  # Oriented boxes object for OBB outputs
    result.show()  # display to screen
    result.save(filename="result.jpg")  # save to disk

 

 

우리가 train 시킨 사과가 없으므로, no detection 이 뜨는 것이 당연하다. 

다른 사과 이미지를 한 번 넣어본다.

train 시킨 사과 이미지와 다르면 된다.

 

 

학습 데이터가 적어서? detect 되지 않는 문제가 발생한다.

val 데이터를 한 번 넣어본다.

 

* val 데이터, train 데이터도 제대로 출력되지 않는다(?)

* yolov8n 모델로 돌리면 제대로 나온다.

 

* 데이터 학습이 제대로 안 되었다는 소리

 

아무튼 yolov8n 모델로 돌려 보았다.

yolov8n.pt

 

test 데이터는 더 있는데, 학습되지 않은 자두와 복숭아에 대해서 사과와 구분을 잘 못하는 모습을 보였다.

그래서, 다른 모델을 적용시켜 보기로 했다.

small, medium, large 까지만 사용해 보기로 했다.

 

yolov8s.pt

small 도 이상하다.

coco dataset 에 복숭아와 자두가 없는게 문제가 큰 것 같다.

 

yolov8m.pt

 

 

yolov8l.pt

 

train 시키지 않은 항목에 대해서 이미 있는것을 바탕으로 추측하고, 사과와 비슷하게 생겼기 때문에 사과로 인지하는 것 같다. 해당 사진은 '자두' 이다.

 

커스텀 데이터셋을 다시 train 시킬 필요가 있겠다.

 

 

 

epoch 를 60으로 늘렸다. train 데이터도 제대로 검출이 안 되기 때문.

아니면 항목이 하나밖에 없어서 그럴 수도 있다.

 

last.pt (epochs=60)

 

best.pt (epochs=60)

 

last 와 best 가 동일한 결과를 보였다.

이제 사과 비슷한 모양은 비슷하게 검출되는 것을 확인 가능하다. (best.pt 사용했다.)

 

 

Full code 도 첨부하도록 하겠다.

필요한 대로 수정하면서 사용하도록 한다.

# --------------------------------------------------
# from ultralytics import YOLO
# model = YOLO('yolov8n.pt')
# model.train(data="C:/Users/admin/Desktop/github/Apple Finder/data.yaml", epochs=60)
# --------------------------------------------------

# from ultralytics import YOLO
# import torch

# model = torch.load("./best.pt")
# print(model)


#----------------------------------------------------
### image predict

from ultralytics import YOLO

# Load a model
model = YOLO("best.pt")  # pretrained YOLOv8n model

# Run batched inference on a list of images
# results = model(["app.jpg"])  # return a list of Results objects
results = model(["apple1.jpg",
                 "apple2.jpg",
                 "apple3.jpg",
                 "peaches1.jpg",
                 "peaches2.jpg",
                 "peaches3.jpg",
                 "plum1.jpg"])

# Process results list
for result in results:
    boxes = result.boxes  # Boxes object for bounding box outputs
    masks = result.masks  # Masks object for segmentation masks outputs
    keypoints = result.keypoints  # Keypoints object for pose outputs
    probs = result.probs  # Probs object for classification outputs
    obb = result.obb  # Oriented boxes object for OBB outputs
    result.show()  # display to screen
    result.save(filename="result.jpg")  # save to disk
#----------------------------------------------------------