객체 탐지 with YOLO: 소개 및 custom dataset 학습하기

기존의 객체 탐지 모델들은 Region Proposal Network(RPN)과 같은 기법을 사용하여, 이미지에서 잠재적인 객체가 있을 법한 여러 영역을 제안한다. 이후, 각 제안된 영역을 분류하고 정확한 바운딩 박스를 예측하는 단계를 거친다. 이 방식은 각 단계에서 별도의 처리가 필요하고, 영역 제안이 많은 경우 계산량이 많아 속도가 느려진다는 단점이 있다. 이 한계를 극복하기 위해 YOLO가 등장하였다.

YOLO는 이미지 전체를 한 번에 처리하여 객체를 탐지하여 매우 빠른 속도와 높은 정확도로 결과를 얻을 수 있다.

YOLO(You Only Look Once) 개요

YOLO란?

YOLO(You Only Look Once)는 한 번의 한 번의 추론으로 전체 이미지에서 객체를 탐지하고 분류하는 작업을 수행하는 모델이다. 연구진은 객체 검출을 하나의 회귀 문제로 정의하여 기존 객체 탐지 모델의 속도와 비용 부담 문제를 해결하였다. 이미지의 픽셀로부터 bounding box의 위치, 클래스 확률을 구하기 까지의 절차를 하나의 회귀 문제로 재정의하였다. 이러한 시스템을 통해 임지 내의 객체 검출을 하나의 파이프라인으로 수행한다. 이미지를 그리드로 나누고, 각 그리드 셀이 예측한 바운딩 박스를 기반으로 객체를 탐지하는 것이다.

YOLO의 주요 특징

빠른 속도: YOLO는 한 번의 네트워크 실행으로 객체를 예측하므로 속도가 빠르다.
높은 정확도: 이미지 내 다양한 객체를 정확하게 탐지하며, 특히 배경과 구분하기 어려운 경우에도 높은 성능을 발휘한다.
다중 객체 탐지: 한 이미지 내에서 여러 개의 객체를 통시에 탐지할 수 있다.

YOLO-Seg란?

YOLO 모델의 확장 버전으로, 객체의 bounding box 탐지뿐만 아니라 객체 분할도 지원한다. 이 기능을 통해 이미지 내 객체의 세부 영역을 더 정확하게 분할할 수 있다.

YOLO-Seg 사용 가이드

YOLOv8-SEG 모델을 전이 학습(Transfer Learning)하여 새로운 데이터셋에 맞춰 커스터마이징하는 과정이다. 전이 학습은 기존에 대규모 데이터셋에서 학습된 모델을 새로운 데이터에 맞게 재학습하는 방법으로, 학습 시간 단축과 성능 향상에 큰 장점이 있다.

전이 학습의 준비 단계

데이터셋 준비 먼저, 새로운 데이터를 준비해야 한다. 데이터는 이미지와 분할 레이블(마스크)로 구성되어 있어야 하며, 다음과 같은 형식으로 구성해야 한다.
- 이미지 파일: 학습할 이미지 파일들을 지정된 폴더에 저장한다.
- 레이블 파일: 각 이미지에 대응하는 분할 마스크를 준비한다. 마스크는 각 객체를 분할하기 위한 라벨 정보가 포함된 파일이어야 한다.’

레이블 데이터는 객체 경계(바운딩 박스 또는 폴리곤)에 대한 정보를 포함하고 있다. 내가 사용한 데이터의 경우 레이블 데이터가 png 파일로 존재했기 때문에, 이를 txt 파일로 변환하는 과정을 거쳤다.

with open(yolo_label_file, "w") as f:
    for contour in contours:
        # Get bounding box
        x, y, w, h = cv2.boundingRect(contour)
        x_center, y_center = (x + w / 2) / label_img.shape[1], (y + h / 2) / label_img.shape[0]
        width, height = w / label_img.shape[1], h / label_img.shape[0]

        # Get segmentation points (relative to image size)
        segmentation_points = [
            f"{pt[0][0] / label_img.shape[1]:.6f} {pt[0][1] / label_img.shape[0]:.6f}"
            for pt in contour
        ]
        segmentation_str = " ".join(segmentation_points)

        # Write to YOLO format: class_id x_center y_center width height segmentation_points
        f.write(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f} {segmentation_str}\n")

환경 설정

전이 학습 과정

데이터셋 구성 파일 작성

names는 각 클래스의 이름, nc는 number of class로 class의 수를 작성한다.

names: ['plant']

nc: 1
train: train 폴더 경로
val: valid 폴더 경로

데이터 폴더 구성은 다음과 같다.

dataset
│
├── train/
│   ├── images/                
│   └── labels/               
│
├── valid/
│   ├── images/                    
│   └── labels/   
│
└── data.yaml

사전 훈련된 모델 로드 및 학습 설정
학습 진행

위 설정을 완료하고 실행하면 된다. YOLOv8은 자동으로 학습 로그를 기록하며, 성능 평가 지표로 mAP(mean Average Precision)과 같은 값들을 출력해준다.

yolo task=segment mode=train model=yolov8s-seg.pt data=data.yaml 경로 epochs=10 imgsz=640

task=segment: 수행할 작업을 지정한다. 옵션에 따라 detect, classify, segment, pose 작업을 선택할 수 있다.
mode=train: 실행 모드를 지정한다. 옵션으로는 train, val, predict 등이 있다.
model=yolov8s-seg.pt: 사용할 모델을 지정한다. (yolov8n.pt, yolov8s.pt, yolov8m.pt, yolov8l.pt, yolov8x.pt 등)
data=data.yaml: 데이터셋 설정 파일(.yaml)의 경로를 지정한다.
epochs=10: 학습 에폭 수를 지정한다.
imgsz=640: 입력 이미지 크기를 지정한다. 기본값은 640이다.
name=runs: 결과 디렉토리의 이름을 지정한다.

학습이 성공적으로 시작되면, 아래와 같이 학습 현황이 출력된다. yolo

그리고 결과 디렉토리에 다음과 같이 log, best 모델 등이 저장된다.

runs
└─segment
   └─train
       │  args.yaml
       │  events.out.tfevents.1731153082.user.17056.0
       │  labels.jpg
       │  labels_correlogram.jpg
       │  results.csv
       │  train_batch0.jpg
       │  train_batch1.jpg
       │  train_batch2.jpg
       │
       └─weights
               best.pt
               last.pt

결과 평가
학습 로그 및 평가

학습이 완료되면, 모델의 성능을 평가하기 위해 검증 데이터셋의 IoU(Intersection over Union)와 mAP를 확인한다.

모델 예측 및 테스트

학습된 모델을 이용해 새로운 이미지에서 객체 분할을 수행할 수 있다.

yolo task=segment mode=predict model=학습 모델 파일 경로 conf=0.25 source=이미지 파일 경로 save=true