YOLO v8과 SAM을 융합한 해충 검출 모델

딥러닝 - 지난 글

• 딥러닝 포스트

참고 링크 (GitHub):
SAM
SAM-annotator
YOLO v8

프로젝트 GitHub: GitHub: YOLOv8 with SAM

YOLO v8이란?

Object Detection, 객체 인식은 이미지 또는 비디오에서 개체를 식별하고 찾는 것과 관련된 컴퓨터 비전 작업이다. Object Detection 기술은 다음 두 가지 질문을 위해 존재한다.

• 이것은 무엇인가? 특정 이미지에서 대상을 식별하기 위함
• 어디에 위치해있는가? 이미지 내에서 개체의 정확한 위치를 설정하기 위함

기존 객체 인식은 다양한 접근 방식으로 시도되면서 데이터 제한 및 모델링 문제를 해결하고자 하였다. 그리고 높은 검출 속도와 정확도를 보이는 단일 알고리즘 실행을 통해 객체를 감지할 수 있는 YOLO 알고리즘이 등장하게 되었다.

YOLO v8은 23년 7월 기준 가장 최근에 나온 YOLO 버전이며 2023 년 1월 Ultralytics 에서 개발되었다. YOLO 모델을 위한 완전히 새로운 리포지토리를 출시하여 개체 감지, 인스턴스 세분화 및 이미지 분류 모델을 train하기 위한 통합 프레임워크로 구축되었다.

자세한 내용은 아래 링크에 잘 설명되어 있으니 참고하면 좋을 것.
참고링크1
참고링크2

YOLO 학습 개요

1. 학습이란?

AI 학습이라는 말이 있다. 말 그대로 학습을 하기 위해서는 문제와 답안지가 필요한 것 처럼, YOLO 모델을 학습하기 위해서 또한 문제와 답안지를 만들어 주어야 한다.

현재 우리는 문제는 가지고 있다. 바로 이미지이다. 우리는 AI가 이미지를 보고 이미지 안에 있는 해충의 종류와 위치 그리고 갯수를 알려주는 답을 얻고 싶은 것이기 때문에 이미지 안에서 해충의 위치를 알려줄 수 있는 답안지를 만들어야 한다. 그리고 YOLO 알고리즘이 이미지와 답안지를 가지고 학습을 하여 만든 하나의 복잡한 풀이과정이 우리가 원하는 모델이 되는 것이다.

2-1. 학습 데이터 준비

YOLO 모델 학습을 위해서는 위에서 언급했듯이 이미지 내에서 어떤 해충이 어디에 위치했는지를 알려주어야 한다. 알려주는 방법으로는 다음과 같은 방법이 있으며 우리는 이에 해당하는 작업을 어노테이션 또는 라벨링이라 할 것이다.

1. Bounding Box
2. Instance Segmentation (polygon)

여기서 우리는 두가지을 한번에 처리할 것이다. Instance Segmentation을 위해 Polygon으로 라벨링한다면 Bounding Box 데이터는 비교적 쉽게 만들 수 있다. 우선 Instance Segmentation을 위한 라벨링 방법은 다음 자료에 자세히 기술하였으니, 참고하면 좋을 것이다.

위 자료를 참고하여 Polygon 라벨 데이터를 만들었다면 확장자 json 형태로 저장 될 것이며, 아래 코드를 참고하여 polygon을 bounding box 데이터로 자동으로 변환 할 수 있다.

학습을 위한 데이터는 총 데이터의 8:2 비율로 학습 및 검증에 사용할 것이다. 다음은 학습에 사용한 데이터 정보이다. 이미지와 라벨 데이터의 갯수는 동일해야만 한다.

진딧물 (aphid)	가루이-약충	가루이-성충

2-2. Roboflow

Roboflow는 컴퓨터 비젼(computer visiom) 기술을 이용해 다양한 어플리케이션을 만들 수 있도록 지원하는 서비스이다. 웹으로 구성되어 쉽게 사용이 가능하고, 주요한 특징은 무료 데이터셋을 제공한다. 사용자가 가지고 있는 데이터를 업로드하여 annotation을 할 수 있다.

roboflow 홈페이지

roboflow에서 프로젝트를 만들고 위에서 작업한 라벨 데이터 (json)과 원본 이미지를 업로드하면, YOLO 학습에 사용할 수 있는 포맷으로 자동 변환해주며 데이터의 증강 또한 가능하기 때문에 부족한 데이터의 갯수를 늘려 학습할 수 있다. 전처리 정보는 아래 사진과 같다 (참고 자료).

증강후 최종적으로 YOLO 학습에 사용한 데이터 정보는 다음과 같다.

라벨 갯수 정보

roboflow에서 모든 세팅을 마친 후 export 하면 아래와 같은 형태의 폴더를 다운받을 수 있으며, 각각 전처리된 이미지, 라벨 데이터와 data.yaml 파일이 포함되어있다.

data.yaml은 YOLO 학습에 사용할 이미지 및 라벨 데이터의 경로를 잡아주는 역할을 하며, 모델에 들어갈 라벨 개수 및 정보 등을 포함한다.

여기까지 왔다면 학습을 시작할 준비가 완료된 것이다.

3-1. YOLO 학습 로깅

Weight & Bias (wandb)는 AI 모델 학습시 진행상황을 서버로 부터 전송받아 대시보드 형태로 출력해주는 사이트로 학습진행상황을 한눈에 보기 좋은 점에 있어서 필수로 사용하는 Tool이다. (참고자료)

YOLO의 경우 wandb를 공식적으로 지원하기 때문에 코드 두줄이면 간단하게 연동할 수 있다.

wandb 설치

pip install wandb

wandb 로그인

wandb login

3-2. YOLO 학습

YOLO는 v8 부터 프레임워크를 지원하기 때문에 ultralytics 라이브러리 하나만 설치하면 간단하게 모델을 학습하고 사용할 수 있다.

Anaconda Prompt에서 CUDAtoolkit, torch, ultralytics를 설치하고 학습 명령어를 입력하면 된다.

학습 파라미터는 다음과 같다. batch는 그래픽카드의 메모리 용량에 따라 설정하면 되며, workers는 기본 8로 설정되어있으나, 그래픽 카드 성능이 부족할 경우 메모리 부족으로 학습이 안될 수 있으니 0 또는 1로 설정하였다.이미지 사이즈는 1,024 x 1,024로 설정하였다.

이를 CLI 상에 다음과 같이 입력하면 학습이 시작된다. (data는 data.yaml이 있는 경로 설정)

yolo detect train data=dataset_3/data.yaml model=yolov8n.pt epochs=300 imgsz=1024 batch=16 workers=0

3-3. YOLO 학습 결과

Confusion Matrix

Train & Validation Loss and Precision

F1-score

진딧물 (aphid)	가루이-약충	가루이-성충

4. YOLO + SAM

SAM (Sagment Anything Model)이란 무엇인지는 (참고 자료)를 읽어보면 알 수 있다.

아래 코드는 학습시킨 YOLO를 구동하여 도출한 Bounding Box 좌표를 SAM에 입력하여 최종적으로 객체를 Instance Segmentation 하는 코드이다. 자세한 내용은 (GitHub: YOLOv8 with SAM 참고)

최종적으로 도출한 결과는 다음과 같다.

진딧물 (aphid)	가루이-약충	가루이-성충