최근에 YOLO라는 것을 알게 되어 나름 검색해 보면서 공부중인데 이 분야를 전공하지 않는 관계로 간략한 사용법 중심으로 살펴보고 있습니다.
기본적으로 제공되는 YOLO의 pre-trained된 파일들이 제공되는데 5개의 (YOLOv8n.pt, YOLOv8s.pt, YOLOv8m.pt, YOLOv8l.pt, YOLOv8x.pt) pt파일이 있으며 nano, small 등 사이즈별로 구별되며 속도 또는 검출 능력별 차이가 있는것 같습니다.
이렇게 미리 훈련된 데이터외에 사용자가 YOLO v8에 custom dataset를 적용하여 데이터가 제대로 검출(detect)되는지 궁금하여 custom dataset를 만들어서 테스트해 보았습니다.
먼저 custom dataset를 만들기 위해서는 대량의 이미지가 필요한데 이러한 이미지들은 구글 검색을 하여 다운로드 받거나 직접 카메라로 획득한 이미지를 사용하여 만들 수 있습니다.
제가 만들 custom dataset는 테스트를 위해 바둑알 이미지로 만들어 볼
예정입니다.
검정돌과 흰돌의 이미지를 30장씩 캡쳐 후 학습시켜 YOLO v8을
사용하여 잘 검출되는지를 테스트해 볼 예정입니다.
먼저 웹캠을 사용하여 샘플용 이미지 30+30장 총 60개의 이미지를 준비합니다.
아래는 제가 이미지 캡쳐를 위해 사용한 코드입니다.
import cv2
import time
count = 0
maxFrames = 30
cap=cv2.VideoCapture(0)
while cpt < maxFrames:
ret, frame = cap.read()
frame=cv2.resize(frame,(640,480))
cv2.imshow("test window", frame)
cv2.imwrite("images/%d.jpg" %count, frame)
time.sleep(0.5)
count += 1
if cv2.waitKey(1)&0xFF==27:
break
cap.release()
cv2.destroyAllWindows()
아래 이미지는 캡쳐된 60장의 바둑알 이미지입니다.
여러장의 이미지에서 실제 검출될 부분을 마킹하는 작업이 필요합니다. 이러한
작업을 쉽게 사용할 수 있도록 제공해 주는 사이트가 있습니다. 아래
사이트( https://universe.roboflow.com )인데 공개된 dataset도 많고 쉽게 custom dataset를 만들수 있습니다. 아래
사이트를 이용하기 위해서는 가입이 필요합니다.
가입 후 아래 화면이 나오는데 왼쪽 상단에 있는 프로젝트를 선택 후 아래쪽의 Create New Project를 클릭하여 새 프로젝트를 작성합니다.
프로젝트 생성창이 나오면 Project Type은 Object Detection을 선택하였고, 그 아래 프로젝트명에 Stone Detection을 detecting 사물 이름으로는 stone으로 적은후 Create Public Project 버튼을 클릭하여 프로젝트를 생성하였습니다.
이렇게 프로젝트를 생성하면 아래와 같이 dataset를 만들 이미지 업로드 화면이 나옵니다.
위 웹캠에서 캡쳐된 60장의 바둑알 이미지를 업로드 한 후 오른쪽 위쪽의 Save and Continue 버튼을 클릭합니다.
왼쪽 메뉴가 Annotate로 변경됩니다.
오른쪽의 메뉴는 이 dataset의 annotation할 팀원을 추가하는 화면 같은데, 아무래도 이 작업이 데이터가 많을수록 좋기 때문에 혼자 하기 보다는 여럿이서 작업을 나눠서 할 경우 팀원을 추가하여 할당합니다. 오른쪽 아래의 Assign Images를 클릭합니다.
총 60장의 이미지의 annotation 할 준비가 되었습니다. 오른쪽 위의 Start Annotating 버튼을 클릭합니다.
클릭하면 아래 Annotation 작업화면으로 이동되는데 제 경우 1번 버튼 Smart Polygon을 선택 후 작업을 하였습니다. 오른쪽 버튼 메뉴에 마우스 커서를 가져다 옮기면 관련 기능을 알 수 있습니다.
Smart Polygon 버튼을 클릭 후 검정 바둑알을 클릭하면 2번과 같이 해당 이미지 부분을 자동으로 선택 해 줍니다.
선택이 되면 3번 팝업창이 나오는데 Finish버튼 또는 Enter키를 눌러 선택 작업을 마칩니다.
선택 작업이 완료되면 Annotation Editor 팝업이 나오는데 class명으로 black_stone이라고 적습니다. 이로써 1개의 이미지의 Annotation 작업이 완료 되었습니다.
5번 화살표를 클릭하여 다음 이미지로 전환 후 같은 작업을 반복합니다.
흰돌의 경우 class명을 white_stone으로 설정합니다.
Annotation 작업이 완료되면 왼쪽 맨 위쪽의 back 버튼을 눌러 메인화면으로 되돌아
옵니다.
아래 이미지의 오른쪽 위쪽의 Add images to Dataset 버튼을
클릭합니다.
아래 이미지의 오른쪽 아래의 Add images 버튼을 클릭하여 train 70%/ valid 20%/test 10%로 이미지가 분배되도록 합니다.
어떤 이미지들이 어떻게 분배되었는지 확인할 수 있습니다.
왼쪽의 Generate 메뉴를 선택하면 아래 이미지와 같이 나오는데 설명을 보면 훈련시간을 줄이고 성능 향상을 위해 이미지를 640x640 정사각형 형태로 변형한다라고 나와 있습니다. 변형된 설정을 변경하지 않고 Continue 버튼을 클릭하였습니다.
마찬가지로 4번 항목도 Continue 버튼을 클릭합니다.
5번 항목의 Generate 버튼을 클릭하면 custom dataset의 생성이 완료됩니다.
이렇게 생성된 custom data를 다운로드 하기 위해 오른쪽 위의 Export Dataset 버튼을 클릭 후 Format을 YOLOv8을 선택합니다.
Format은 YOLOv8을 선택 후 Google Colab에서 사용하기 위해 아래 'show download code'를 선택 후 Continue 버튼을 클릭합니다.
다운로드 코드창이 나오는데 Colab에서 사용될 Raw URL을 선택 후 아래 붉은색 주소를 복사해 놓습니다. 이 주소를 Colab에서 사용하면 위에서 작업된 dataset를 불러 올 수 있습니다.
dataset관련 작업은 모두 끝났습니다. 다음은 Google Colab에서 roboflow에서 작업된 Stone dataset를 불러와 학습시키는 과정입니다.
새 프로젝트를 생성 후 메뉴에서 런타임 유형을 T4 GPU로 변경합니다.
Training
작업이 워낙 오래 걸리는 작업이라 기본 CPU로 작업하면 너무 오래 걸립니다.
아래 명령어를 사용하여 Colab에서 학습시킵니다.
Colab으로 roboflow 생성한 dataset를 stone.zip으로 다운로드 받습니다.
주소는 roboflow에서 Export시 생성된 URL 입니다.
!wget -O stone.zip <여기에 roboflow에서 만든 stone dataset 주소를 적어 줍니다.>다운로드 받은 dataset인 stone.zip 파일의 압축을 해제합니다.
import zipfile
with zipfile.ZipFile('/content/stone.zip') as target_file:
target_file.extractall('/content/stone')
PyYAML을 설치합니다.
!pip install PyYAML기존 YAML 파일 열어서 확인해 봅니다.
!cat /content/stone/data.yaml기존 YAML에 새로운 Data를 덮어씌웁니다.
import yaml
data = {'train' : '/content/stone/train/images',
'val' : '/content/stone/valid/images',
'test' : '/content/stone/test/images',
'names' : ['black_stone', 'white_stone'],
'nc' : 2}
with open('/content/stone/data.yaml','w') as f:
yaml.dump(data, f)
with open('/content/stone/data.yaml','r') as f:
Fire_yaml = yaml.safe_load(f)
display(Fire_yaml)
Colab에 ultralytics 설치합니다.
!pip install ultralyticsultralytics의 버전을 확인해 봅니다.
import ultralytics ultralytics.checks()YOLO에서 yolov8n.pt를 Colab에 받아옵니다.
from ultralytics import YOLO
model = YOLO('yolov8n.pt')
다운받은 yolov8n.pt의 Class 확인해 봅니다.
print(type(model.names), len(model.names)) print(model.names)새로운 stone class로 훈련시킵니다.
model.train(data='/content/stone/data.yaml', epochs=100, patience=30, batch=32, imgsz=416)훈련시킨 yolov8n.pt의 Class가 black_stone과 while_stone으로 나오는지 클래스를 확인 해 봅니다.
print(type(model.names), len(model.names)) print(model.names)훈련된 best.pt를 다운로드 받습니다.
from google.colab import files
files.download('/content/runs/detect/train/weights/best.pt')
YOLO v8에서 다운로드 받은 best.pt를 사용하여 custom dataset이 정상적으로
동작되는지 확인해 봅니다.
yolo predict model=best.pt source=0 show=True
위 명령어를 사용하여 바둑알 검출을 해 봅니다.
아래 그림과 같이 검정돌과 흰돌이 정확히 검출됩니다.
이런식으로 roboflow에서 자신만의 custom dataset를 생성하여 YOLO에 적용 해 보았습니다.