YOLOv5를 iOS에서 사용하는 방법입니다

YOLOv5는 객체 탐지 (Object Detection) 머신러닝 모델입니다.

iPhone이나 iPad에서 사용할 수 있도록 CoreML 형식으로 변환합니다.

변환한 모델은 그대로는 사용할 수 없습니다

CoreMLTools로 심플하게 변환한 YOLO 모델의 출력은 클래스별로 방대한 양의 박스(Box)이므로, 그대로는 사용할 수 없습니다.

디코딩 (Decode)용 레이어와 Non-Max Suppression (NMS) 레이어를 추가하여, 신뢰도가 높은 박스의 좌표를 압축할 필요가 있습니다.

이러한 처리를 함으로써 iOS의 Vision 프레임워크에서 다룰 수 있게 되며, 또한 서두의 이미지와 같은 간단한 파일 미리보기 기능을 사용할 수 있게 됩니다.

iOS에서 사용할 수 있도록 모델을 편집하는 스크립트가 여기 있습니다

변환 완료된 모델 (COCO 데이터셋)은 여기 있습니다.

독자적인 데이터셋으로 트레이닝한 YOLOv5 모델도 이 변환 스크립트를 통해 CoreML 형식으로 만들 수 있습니다.

변환 코드 해설

YOLOv5 리포지토리 (Repository) 내에 있는 익스포트 (Export) 코드를 이용하여, PyTorch 모델을 CoreML 모델로 변환합니다.

python export.py --weights yolo5s.pt --train --include "coreml"

디코딩 레이어를 정의합니다.

# Just run to define the decode function
import torch
# classLabels = [f"label{i}" for i in range(80)]
...

Non-Max Suppression을 정의합니다.

# Just run to define the NMS function
def createNmsModelSpec(nnSpec):
'''
...

CoreML 모델에 디코딩 레이어와 Non-Max Suppression을 추가합니다.

# run the functions to add decode layer and NMS to the model.
addExportLayerToCoreml(builder)
msSpec = createNmsModelSpec(builder.spec)
...

이렇게 하면 Vision에서 사용할 수 있는 CoreML 모델이 저장됩니다.

iOS에서의 사용 방법

디코딩과 NMS는 다음 리포지토리와 서적을 참조하고 있습니다.

🐣

프리랜서 엔지니어입니다.

업무 관련 상담은 이쪽으로

rockyshikoku@gmail.com

Core ML이나 ARKit을 사용한 앱을 만들고 있습니다.

머신러닝 / AR 관련 정보를 발신하고 있습니다.

Discussion