Yolov5 커스텀 데이터 학습 Object detection

YOLOv5는 단일 단계 방식의 객체 감지 딥 러닝 기법으로, 빠른 속도와 실시간 객체 인식 능력을 가지며 Fast R-CNN과 같은 기존의 실시간 객체 탐지 모델들과 비교하여 높은 성능을 보임.

단일 단계 방식은 입력 이미지를 한 번의 전방향 패스로 처리하여 객체의 위치와 클래스를 예측하는 방법이므로 기존 R-CNN 계열의 모델들과 달리 여러 단계의 처리를 거치지 않아도 되어 속도가 훨씬 빠르다는 특징을 가짐

Fast R-CNN과 같은 다른 모델들은 여러 단계로 나누어 객체 탐지를 수행하는데, 먼저 후보 영역을 추출한 다음, 각 후보 영역을 신경망에 넣어 객체의 위치와 클래스를 예측하지만 YOLOv5는 이미지를 그대로 신경망에 통과시켜 한 번의 추론으로 객체 탐지를 수행하므로 빠른 속도와 실시간 감지가 가능함. 

arxiv.org/pdf/1506.02640.pdf

Yolov5는 COCO 데이터셋으로 학습되었지만 커스텀 데이터셋을 활용하여 파인튜닝 또한 가능함.

다양한 데이터셋과 직접 레이블링 할 수 있는 roboflow를 활용하였음.

https://universe.roboflow.com/

Roboflow Universe: Open Source Computer Vision Community

 

Colab을 이용하였고 런타임에서 gpu사용 활성화.

야구방망이를 인식하는 모델을 만들기 위해 roboflow에 야구방망이에 대해 레이블링이 되어있는 데이터셋 3개를 이용하여 train과 validation 데이터셋으로 나눔.

#Yolov5 다운로드
!git clone https://github.com/ultralytics/yolov5
%cd yolov5
%pip install -qr requirements.txt

# 데이터셋 다운로드 및 세팅
!mkdir /content/dataset
%cd /content/dataset
!curl -s -L "https://universe.roboflow.com/ds/RCHmjfUtI6?key=6FLwSuWNt5" > roboflow.zip; unzip -q roboflow.zip; rm roboflow.zip
!mkdir /content/dataset2
%cd /content/dataset2
!curl -s -L "https://universe.roboflow.com/ds/Gtk9PToLwg?key=bXvrVowYc4" > roboflow.zip; unzip -q roboflow.zip; rm roboflow.zip
!mkdir /content/dataset3
%cd /content/dataset3
!curl -s -L "https://universe.roboflow.com/ds/HHNf0N05mb?key=XVsRZj3MSE" > roboflow.zip; unzip -q roboflow.zip; rm roboflow.zip
!mv /content/dataset/train/images/* /content/dataset3/valid/images
!mv /content/dataset2/train/images/* /content/dataset3/train/images
!mv /content/dataset/train/labels/* /content/dataset3/valid/labels
!mv /content/dataset2/train/labels/* /content/dataset3/train/labels
!mv /content/dataset3/test/images/* /content/dataset3/train/images
!mv /content/dataset3/test/labels/* /content/dataset3/train/labels
%cd /content/yolov5
!rm -rf /content/dataset /content/dataset2

데이터셋의 개수를 확인하기 위한 코드

# train 폴더의 이미지 및 레이블 파일 개수 확인
!echo "Number of images in train folder:" $(ls -1 /content/dataset3/train/images | wc -l)
!echo "Number of label files in train folder:" $(ls -1 /content/dataset3/train/labels | wc -l)

# valid 폴더의 이미지 및 레이블 파일 개수 확인
!echo "Number of images in valid folder:" $(ls -1 /content/dataset3/valid/images | wc -l)
!echo "Number of label files in valid folder:" $(ls -1 /content/dataset3/valid/labels | wc -l)

#Number of images in train folder: 3690
#Number of label files in train folder: 3690
#Number of images in valid folder: 1269
#Number of label files in valid folder: 1269

data.yaml 파일에서 train과 validation 데이터셋의 경로를 지정해주고 nc는 number class와 class name 또한 수정함.

#yaml파일 수정
!mv /content/dataset3 /content/dataset
with open('/content/dataset/data.yaml', 'r') as file:
    yaml_content = file.read()
print(yaml_content)
new_yaml_content = """
train: /content/dataset/train/images
val: /content/dataset/valid/images

nc: 1
names: ['bat']
"""
with open('/content/dataset/data.yaml', 'w') as file:
    file.write(new_yaml_content)

 

Yolov5는 모델마다 추론속도 및 정확도가 다름. 추론속도가 빠르다면 정화도는 다소 떨어질 수 있고, 정확도가 높다면 추론 속도가 떨어지는 관계고, 필자는 중간모델인 medium을 사용

 

#yolov5m 다운로드 및 훈련
!wget https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5m.pt
!python /content/yolov5/train.py --img-size 640 --batch-size 16 --epochs 50 --data /content/dataset/data.yaml --weights /content/yolov5/yolov5m.pt

--source는 분석하고 싶은 사진이나 영상 --weight는 가중치 경로를 지정

#테스트
!python /content/yolov5/detect.py --source "test.mp4" --weight /content/yolov5/runs/train/exp/weights/best.pt

젯슨나노로 Yolov5를 통해 웹캠과 연결하여 실시간 감지도 가능함.