Yolov5 데이터셋 학습

장성환·2024년 6월 1일

우리의 프로젝트는 3구 당구에서 당구대와 당구공을 인식해서 각 좌표를 구해주는 프로그램을 만드는 것이다.

데이터셋 준비

필자는 Roboflow라는 사이트에서 라벨링한 데이터셋을 pytorch yolov5용으로 다운받았다.

RoboFlow에 대한 사용법은 검색하면 나오니 검색 고고

데이터셋을 다운받으면 아래와 같은 파일들이 있을것이다.

여기서 조금의 수정이 필요한 부분이 있다.
data.yaml을 열어서,

위와 같이 train, val, test의 폴더 경로를 절대경로로 수정해주어야 한다.
그대로 실행하면 오류뜨면서 막힌다.

수정했으면 학습시킬 차례다.

cd yolov5

명령어로 yolov5경로로 이동후

C:\Users\Jang\yolov5>python train.py --img 640 --batch 20 --epochs 30 --data datasets/data.yaml

여기서 img 는 학습하는 이미지의 크기,
batch는 몇 번마다 가중치를 업데이트 할 건지(16이라면 16번째마다 가중치가 업데이트 된다.)
epoch은 학습의 반복을 의미한다.
만약 이미지가 100개이고 batch가 20, epoch이 30이면 100/20 * 30 = 150번의 학습을 반복하게 되는것이다.

batch와 epoch가 높으면 높을수록 학습량이 많아지므로 학습 속도가 느려진다.

아무튼 위의 명령어를 실행하면

이런식으로 학습을 진행한다.

학습이 끝나면

Results saved to runs\train\exp7

학습된 모델이 어디에 저장되어있는지 알려준다.

마지막으로 학습을 시켰으면 학습이 제대로 되었나 확인해 봐야지.

python detect.py --source ./images/test.png --weights runs\train\exp7\weights\best.pt --conf 0.4

--source 뒤에 객체를 찾을 이미지 경로, --weights 뒤에 아까 저장된 데이터 경로의 weights폴더 안에 best.pt
--conf는 정확도 설정, 1에 가까울수록 깐깐하게 평가한다.


결과가 저장된 사진이 맨 아래 표시된 경로에 저장된다.

이처럼 학습이 완료되어 객체를 잘 탐지 하는 모습을 볼 수 있다.
탐지가 잘 안된다면 epoch이나 batch의 값에 변화를 주거나, 학습 데이터 자체를 늘려보자

profile
홍익대학교 게임소프트웨어전공, 산업공학전공

0개의 댓글