[Computer Vision Study] Yolo 커스텀 모델 만들기

Dodoy·2024년 5월 25일

CV

목록 보기
1/2
post-thumbnail

오늘은 kaggle에서 알파카 이미지 데이터셋을 다운받아 Roboflow에서 Annotation을 한 후
학습된 모델을 이용하여 알파카를 인식하는 작업을 수행해보았다.

데이터셋 다운로드

우선 적당한 이미지 데이터셋을 찾다가 kaggle에서 142장의 알파카 이미지를 다운받았다.

Roboflow

Annotation

Roboflow에서 다운받은 이미지 데이터셋을 업로드하고 142장의 알파카 사진을 bounding box로 annotation했다. (손목 끊어지는 줄 알았음..)

Generate

Train : Valid : Test = 7 : 2 : 1
Augmentation은 flip, rotation, crop 으로 뻥튀기시켰다.

Data Export

모델은 Yolov8로 진행하였다.

Training

https://github.com/Kimdoyeon123/Yolo_study01

Yolov8 기반으로 학습한 pt파일을 이용하여 알파카 사진을 인식하는 작업을 최종적으로 수행했다.

import ultralytics
from ultralytics import YOLO

model = YOLO('yolov8s.pt')

model.train(data='Alpaca_detection.v2i.yolov8/data.yaml' , epochs=20)

epoch 20으로 학습을 진행하였다.


import cv2
from ultralytics import YOLO

model = YOLO('runs/detect/train/weights/best.pt')
results = model.predict(source="C:/Users/user/Downloads/al.jpg", show=True, save=True, conf = 0.02, show_boxes=True)

바운더리 박스로 알파카를 인식하고 정확도를 표시하였다.

Yolo 모델과 알파카 커스텀 모델

Yolo에 기본적으로 탑재된 객체인식 모델은 알파카를 양과 새로 인식을 했다.

그렇지만 내가 학습시킨 알파카 모델은 configure 를 너무 낮게 해서인지 알파카들 사이의 틈과 알파카의 몸통과 머리를 따로따로 인식을 해서 알파카를 7마리로 인식했다.

더 많은 데이터와 정확한 라벨링이 필요할 것 같다.

profile
Ai 공부중

0개의 댓글