[Object Detection] 누구나 쉽게 따라할 수 있는 YOLOv5 모델 학습하기 (커스텀 데이터)

Object Detection

Object Detection(객체 검출)이란 이미지에서 검출하고 싶은 객체의 위치(바운딩 박스)를 찾고 객체의 카테고리를 분류하는 작업을 의미한다. 객체 검출하면 역시 많은 분들이 YOLO 모델을 떠올리는 것 같아서, YOLO 모델을 학습하고 테스트하는 튜토리얼을 준비해 봤다.

Object Detection 예시

요즘은 컴공과나 AI 관련 공부를 하시는 학생이 아니더라도 본인의 학업에 객체 검출 기술을 적용시켜보려 하는 사람들이 많다. 하지만 코딩을 할 줄 모르면 간단한 코드일지라도 따라하다보면 에러가 나고 어떻게 해결해야할지 모르겠는 상황이 많이 연출된다. 때문에 코딩을 할 줄 모르는 사람도 쉽게 따라할 수 있는 YOLO 학습 튜토리얼을 준비했으니 필요한 분들께 도움이 되었으면 좋겠다.

 

- 학습 코드 레퍼지토리 : https://github.com/ldj7672/Deep-Learning-Tutorials/tree/main/YOLOv5
(위 레퍼지토리에서 "YOLOv5_학습_커스텀데이터셋.ipynb" 파일 사용)

* 코드를 다룰줄 안다면 VS code 같은 IDE로 개발환경을 구축해서 모델을 학습시키는 방법이 훨씬 간편하다

데이터셋 준비

데이터셋 선정 및 다운로드

 

• Roboflow 퍼블릭 데이터셋 :https://public.roboflow.com/object-detection

궁극적으로 데이터셋은 어떤 것을 사용하셔도 되지만 간편하게 roboflow 라는 사이트에서 object detection 을 위한 여러가지 퍼블릭 데이터셋을 다운받아 사용할 수 있다.

그 중에서 나는 드론으로 부둣가를 촬영하여 자동차, 배, 부두 등을 검출하기 위해 만들어진 'Aerial Maritime Drone Dataset' 이라는 데이터셋을 선택했다. 위 roboflow 사이트에 접속해서 데이터셋을 검색하고 다운받을 수 있다.

다운로드 받을 때 데이터셋 포맷은 TXT - YOLOv5 pytorch 포맷으로 다운받으시면 된다. 해당 포맷은 이미지 어노테이션에 대한 포맷인데 코드와 어노테이션 포맷이 서로 일치해야 에러가 발생하지 않는다.

만약 본인이 가진 데이터셋으로 학습하기 위해선 퍼블릭 데이터셋과 동일한 포맷으로 annotation(정답 라벨(바운딩 박스, 카테고리))을 만들어서 적용시켜줘야 한다. 어떻게 해야 할지 모르겠으면 용량이 작은 데이터셋을 하나 다운받아서 동일한 포맷과 동일한 디렉토리 구성으로 만들어주면 된다.

Roboflow 퍼블릭 데이터셋

Aerial Maritime Drone Dataset

데이터셋 포맷

 

다운받은 데이터셋 확인

다운받은 압축 파일을 해제하면 아래와 같이 학습/평가 데이터와 정답 정보가 들어있는 test, train, valid 폴더와 데이터셋에 대한 정보가 들어있는 data.yaml 파일이 있다. yaml 파일은 코드 에디터나 텍스트 편집기로 열어보시면 아래처럼 확인이 가능한데 나중에 yaml 파일에서 train, val 경로를 변경해줘야 한다.(밑에서 설명)

좌 : 다운 받은 데이터셋 폴더, 우 : yaml 파일

 

데이터셋, 학습 코드 → 구글 드라이브 업로드

이제 다운 받은 데이터셋 폴더를 구글 드라이브에 업로드해야 한다. YOLO 모델을 구글 코랩이라는 환경에서 학습시킬 예정이기 때문에 데이터셋을 구글 드라이브에 업로드하는 것

구글 드라이브의 '내 드라이브'에 'dataset' 이라는 폴더를 하나 만들어 주고 그 안에 다운받은 데이터셋을 업로드 시킨다. (폴더 이름에 점(.)이나 각종 특수 문자가 있는 경우 코드 실행 시 문제가 되므로 이름을 바꾸어 줘야 한다)

그리고 제공한 'YOLOv5_학습_커스텀데이터셋.ipynb' 파일또한 '내 드라이브' 경로에 업로드한다.

구글드라이브 - 내 드라이브에 dataset 폴더 생성, YOLOv5_학습_커스텀데이터셋.ipynb 파일 업로드

내 드라이브/dataset에 데이터셋 업로드

이제 데이터셋과 학습 코드준비가 끝이 났다.

YOLOv5 모델 학습 및 테스트 - 구글 코랩 사용

앞서 얘기한대로 YOLOv5 모델은 구글 코랩이라는 환경에서 학습시키려 한다. 구글 코랩을 사용 시 개발 환경을 따로 세팅해줄 필요가없어서 예상치 못한 에러를 만날 확률이 적기 때문에 코딩을 할 줄 모르는 분들도 쉽게 따라할 수 있다는 장점이 있다.

 

1. 구글 드라이브 마운트

 

• 구글 코랩과 구글 드라이브를 연결시켜주기 위해 마운트를 필요 (아래 화면에서 재생표시를 눌러서 진행)
• 빠른 학습을 위해 GPU를 사용 (런타임 - 런타임 유형 변경 - 하드웨어 가속기 → GPU)

 

 

2. YOLOv5 레퍼지토리 다운로드

 

YOLOv5 를 사용하기 위해 깃허브에 올려진 YOLOv5 레퍼지토리에서 코드들을 다운로드. 이는 한번만 실행해주면 구글 드라이브 내 드라이브에 yolov5가 다운로드된다. 두 번째 학습부터는 건너뛰어도 되는 부분이고, 패키지 다운로드 및 임포트 하는 부분은 매번 실행해줘야 한다.

 

3. 데이터셋 다운로드, yaml 파일 수정

데이터셋은 이미 다운로드 받고 구글 드라이브에 업로드시켜서 아래와같이 데이터셋과 yaml 파일을 확인할 수 있다. 여기서 데이터셋의 경로를 담고 있는 yaml 파일을 조금 수정해줘야 학습이 가능하다. 아래 화면에서 data.yaml 을 클릭하셔서 직접 수정해도 되지만, 코드로 yaml 파일을 수정하도록 구현해두었다.

구글 코랩에서 구글 드라이브 파일 확인

우선 아래와 같이 data_dir 에는 데이터셋의 경로를 data_yaml 변수에 yaml 파일의 경로를 지정해 준다. 그리고 나서 아래 코드를 수행하면 yaml 파일의 정보가 나오는데, 우리는 데이터셋을 구글 드라이브에 올렸기 때문에 여기서 train, val 의 경로를 수정해줘야 한다.

코드를 실행하여 yaml 파일 정보를 확인하고, film['train'], film['val'] 부분을 수정하여 yaml 파일을 수정(구글 드라이브에서의 데이터셋 경로)해주고 변경된 정보를 확인한다.

기존 yaml 파일의 정보 확인

yaml 파일 정보 수정

 

 

4. YOLOv5 모델 학습

이제 아래 학습 파라미터들을 지정해주고 모델 학습을 시작하면 된다. 지정한 epoch 만큼 학습을 진행하고 나면 학습이 종료되고 'Results saved to runs/train/exp*' 이라는 문구가 뜨는데 이 경로에 학습한 정보가 저장되었다는 뜻이다. exp 몇번에 저장되었는지 확인 필요.

YOLOv5 모델 하이퍼 파라미터 및 학습 실행 코드

학습 종료 후 저장된 경로 확인 필수 !

5. 텐서보드 시작하기

텐서보드는 학습시키는 딥러닝 모델이 학습되는 과정을 여러 가지 metric으로 기록하여 시각적으로 볼 수 있는 도구이다.  앞서서 확인한 exp 번호를 보고 왼쪽 아래에서 원하는 학습 정보만 필터링해서 볼 수 있다.

6. 학습한 모델 테스트

이제 학습한 모델을 테스트 데이터셋에서 테스트를 해봐야 한다. 앞서 확인한 exp 번호를 train_exp_num에 저장해두고 테스트를 실행한다. 테스트 또한 끝나고 나면 테스트 정보가 저장된 경로가 나오는데 이 exp 번호 또한 확인해둬야 한다.

학습 exp 번호를 지정

테스트 정보가 저장된 exp 번호 확인

확인한 테스트 exp 번호를 text_exp_num에 저장하고 테스트 결과 확인 코드를 실행하면 아래 그림과 같이 이미지에서 객체를 검출한 결과를 확인할 수 있다.

테스트 결과 확인 코드

테스트 결과 예시

드디어 테스트 결과를 이미지로 확인이 가능하다. 결과가 잘 나오지 않았다면... 학습이 어디선가 잘 못된 것이고 코드를 처음부터 잘 확인해햐 한다. 코드가 돌아는 갔는데 학습이 잘 안되는 경우가 가장 막막한 경우인데 데이터로더가 뱉어주는 이미지 데이터와 레이블 데이터를 확인해서 학습에 데이터가 잘 들어가는지 부터 확인하는 것을 추천한다.

 

7. 베스트 모델 저장

마지막으로 학습 중 결과가 가장 좋은 베스트 모델을 저장할 수 있고 저장한 모델은 언제든 불러와서 사용할 수 있게 된다.

 

---

최근에 핫한 이슈인 생성형 AI의 트렌드를 따라가기 위해서는 지속적인 관심이 필요한데요. 여러 산업에서의 유즈 케이스를 잘 파악하는 것도 굉장히 중요합니다. 해서 이번에 생성형 AI 관련 세미나를 소개해 드리려합니다. 서울 코엑스에서 데이터 인텔리전스 플랫폼과 생성형 AI 및 여러 유즈 케이스를 소개하는 세미나이니 많은 관심 부탁드립니다.

 

- https://events.databricks.com/KoreaDIDays2024