YOLO v5를 이용하여 Custom data 인식하기

YOLO v5

영상인식에서 자주 쓰이는 yolo(you only look once) 에서 다섯번째 버젼이라고 할 수 있습니다.
하지만 이전 v4 까지와는 다른 제작자가 만든 것이므로 해외에서는 v5의 이름이 적절한지 논란이 있다고 합니다.
이번 글은 아래 링크를 참고하여 Lego GBC 모듈에서 공을 학습시키는 과정입니다.

 

라벨링 과정

v5에서 학습을 시키기위해 이미지와 객체의 위치가 라벨링 된 파일이 있어야 합니다. 

이 때, 웹사이트에서 구하는 방법이 있는데 저는 구하기 어려운 data 이므로 영상에서 이미지를 추출하는 방식을 사용했습니다.

https://deftkang.tistory.com/182

 

이미지 추출 코드

중요한 점은 사이즈를 416, 416으로 하니 화질이 너무 깨져 거의 쓸 수 없었습니다. 따라서, 해상도를 1920, 1080 정도로 설정하고 따로 resize하는 코드를 해야 할 것 같습니다.

 

 

다음과 같이 이미지를 얻을 수 있는데 저는 첫번째 학습이 실패하여 다시 이미지를 정리하여 98장 정도가 됩니다.

 

Annotation(LabelImg 사용)

라벨링에 있어 박스를 그리는 중요한 프로그램

github.com/tzutalin/labelImg

tzutalin/labelImg

 

 

다운받고 저는 Anaconda를 깔아놓았기 때문에 anaconda prompt 를 사용하여 다운 받은 경로로 이동하고 위처럼 입력하면 아래와 같은 프로그램이 나타나게 됩니다.

LabelImg

open dir 사용하여 사진이 들어간 파일을 선택해주고

change save dir 을 label 이 들어갈 파일을 설정해주면 됩니다. 저는 images, labels 파일을 만들어 저장했습니다.

위처럼 박스를 그리고 저장하면 

 

위처럼 라벨링이 된 txt 파일을 얻을 수 있습니다. 이제 모든 이미지에 해당 작업을 수행해 주면 됩니다.

* 중요한점

라벨링을 하기 전 labelimg-master 폴더로 가셔서 data에 predifined classes 를 본인에 맞게 수정을 꼭 해주셔야 합니다.위 사진에서 보다시피 새로운 label이 15, 16으로 설정되는 것을 볼 수 있는데 이러면 yolo v5를 사용할 때 헷갈리게 되고 본인의 label이 0부터 시작하게 만들어주세요!

 

colab 에서 학습하기

이제 학습할 data를 설정했다면 export 파일을 만들고 yaml 파일을 만들어줍니다(동영상 참고)

 

 

코랩을 실행하고

 

마운트

 

 

YOLOv5 다운

 

필요한 패키지 다운

 

 

 

yaml 파일확인(저는 2개의 클래스이고 각각 B(basketball), S(soccerball)

 

glob을 사용해 경로를 list로 받아옵니다

 

98개의 이미지를 9:1로 split (data가 적어 9:1로 하였습니다)

 

 

위 경로가 담긴 리스트를 txt 파일로 적어 생성합니다.

 

이제 yaml 파일에 위의 경로들을 적어주고

 

마지막 단계인 train.py를 실행해주면 됩니다

%cd /content/gdrive/MyDrive/yolov5/
!python train.py --img 416 --batch 16 --epochs 50 --data /content/gdrive/MyDrive/GBC_balls/GBC_balls.yaml --cfg ./models/yolov5s.yaml --weights yolov5s.pt --name GBC_yolov5s_results

 

여기서 yolov5는 모델이 s, m, l, x 순으로 나누어져 있습니다. 순서대로 모델을 학습하는데 오래걸리는 대신 정확성이 좋다고 합니다.

하지만 저는 똑같은 데이터를 s, x 한번씩 수행하니 x 모델은 객체를 인식하지 못하는 문제가 발생합니다. 데이터가 너무 적어서 그런지 잘 모르겠습니다ㅜ

 

결과물

인식은 하고 있지만

 

빈 바구니를 인식하는 등 인식율이 많이 떨어집니다.

 

따라서, 데이터를 많이 늘려보거나 바구니가 있는 데이터는 없애거나

혹은 다른 후처리를 생각해봐야 할 것 같습니다