Yolo11 커스텀 데이터 학습하기

https://github.com/ultralytics/ultralytics

GitHub - ultralytics/ultralytics: Ultralytics YOLO11 🚀

 

Yolo11이란?

• Yolo11은 ultralytics 에서 컴퓨터 비전과 AI에 대한 연구를 기반으로 만들어진 모델
• 객체감지,추적,인스턴스분할,이미지분류,포즈추정 작업에 탁월

출처 : https://github.com/ultralytics/ultralytics

기존 Yolo 모델과의 COCO mAP (객체 탐지의 정밀도) 성능 비교를 살펴 보면 Yolo11은 이전 버전 및 다른 객체 탐지 모델 대비 성능과 속도 면에서 뛰어난 우위를 점하고 있는것으로 나타난다.

여기서 모델을 살펴 보면 다음과 같다.

• 11n(Nano) : 가장 가벼운 모델로 연산량과 메모리 사용이 매우 적어 저사양 장치에서 실시간 객체 탐지가 필요한 경우 적합
• 11s(Small) : Nano 보다 약간 크고 더 높은 정확도를 제공하며 모바일장치나 임베디디 시스템에서 효율적으로 작동
• 11m(Medium) : 균형 잡힌 성능을 제공하며 대부분의 일반적인 응용 프로그램에 적합
• 11l(Large)  : 더 많은 연산 자원이 필요하지만,높은 정밀도를 제공
• 11x(Extra Large) : 가장 강력한 모델로 최고 수준의 정확도를 제공, 처리 속도 보다는 정밀도가 중요한 고급 애플리케이션(예:의료 이미지 분석)에 적합

또한 전처리 모델을 사용하여 학습데이터인 yaml 파일만 바꾸어 주면 동일한 파이썬 API를 이용해서 다양한 작업을 할 수 있다.

• Detect 모델 : YOLO11n 과 같이 기본 모델
• Segment 모델 : YOLO11n-seg 와 같이 seg가 붙은 모델
• Classify 모델 : YOLO11n-cls 와 같이 cls 가 붙은 모델
• Pose 모델 : YOLO11n-pose 와 같이 pose 가 붙은 모델
• OBB 모델 : YOLO11n-obb 와 같이 obb 가 붙은 모델

 

전처리 모델 사용 훈련 예

from ultralytics import YOLO

# Load a model
model = YOLO("yolo11n.pt")  # load a pretrained model (recommended for training)

# Train the model with 2 GPUs
results = model.train(data="coco8.yaml", epochs=100, imgsz=640, device=[0, 1])

위의 예는 detect 모델을 훈련하는 예제인데 이것을 Segment 모델을 훈련한다고 하면 yolo11n.pt -> yolo11n-seg.pt 모델로 변경 되면 된다.

또한 model.train 부분의 학습데이터가 들어 있는 data 부분의 해당 yaml file 로 변경한다.

각각의 모델에 따라 이 두개의 부분만 변경이 되고 나머지는 모두 동일하다.

 

Yolo11 커스텀 데이터 학습 순서

1. 데이터 준비
2. 데이터 로딩
3. YAML File 생성
4. Yolo11 인스톨
5. 모델 훈련
6. 예측

1. 데이터 준비

Custom Data로 Yolo11 모델을 학습하는 경우 Image/Annotaion 으로 이루어진 Data를 준비해야 한다.

이러한 Custom Data는 Roboflow(https://universe.roboflow.com/) 에서 제공하는 Custom Data를 이용하거나 혹은 직접 Custom Data 를 구축해야 된다.(Custom Data 를 구축하는 경우 이미지 데이터와 정답데이터는 파일이름이 동일해야 하며 정답데이터(즉 Annotaion)는 파일명이 .txt 여야 한다.

이번 시간에는 CCTV로 화재가 발생했을 때 감지하여 알람을 발생하는 모델을 만들어 보도록 하자.

데이터셋 위치 : https://universe.roboflow.com/5w20vzqobtxjjhtjq6kad9ubrm33/wildfire-smoke/dataset/1

Wildfire Smoke Object Detection Dataset (v1, raw) by AI For Mankind

YOLOv11 을 선택한다.

Download dataset 을 선택한다.

Show download code 를 선택한다.

위와 같이 Raw URL 을 제공해 준다.

 

2. 데이터 로딩

1. 먼저 colab의 런타임을 GPU로 설정한다.

2. 다음의 명령어로 데이터를 다운로드 받는다.

!mkdir dataset
import os
os.chdir('./dataset') #dataset 폴더로 작업 디렉토리 이동

!curl -L "https://universe.roboflow.com/ds/faAuJ1ZGRm?key=5Uzhlu3Ojs" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

3. 디렉토리 내에 파일이 생성 된 것을 확인한다.

3. data.yaml 파일 살펴 보기

!cat data.yaml

train,val,test 위치가 상위 디렉토리로 설정 되어 있는 것을 알 수 있다. 이 부분을 데이터셋에 맞춰서 수정해 주어야 한다.

 

3.YAML File 생성

train,val,test 를 이미지 셋이 있는 디렉토리 위치로 변경을 해 주어야 한다.

1. data.yaml 을 더블클릭하면 우측에 파일 에디터 형식으로 에디팅을 할 수 있는 창이 뜬다.

2. 파일 탐색기에서 다운로드 받은 파일 위치를 우클릭하여 경로복사를 한다.

3. 다음으로 data.yaml 의 파일 디렉토리를 다음과 같이 전체경로 형식으로 변경한다.

 

4. Yolo11 인스톨

!pip install ultralytics

5. 모델 훈련

1. 사전학습 모델 불러오기

import ultralytics
from ultralytics import YOLO

model = YOLO('yolo11n.pt') #사전 학습된 yolo11n 모델을 로드한다.

여기서는 detect 모델을 사용했다.

2. 사전학습된 모델의 정보 살펴 보기

print(type(model.names),len(model.names))
print(model.names)

dict 형태의 80개 class으로 사전학습되어 있는 모델을 확인할 수 있다.

3. 모델 훈련

model.train(data='/content/dataset/data.yaml',epochs=30,patience=5,imgsz=416)

data 에 우리가 만든 yaml 위치를 작성하고 30번 학습을 시켜 보자.

• epochs : 30번 학습
• patience : 조기종료 기준으로 5에포크 동안 개선되지 않으면 종료
• imgsz : 모델에 입력되는 이미지의 크기- 여기서는 416*416 픽셀로 리사이징하여 사용

4. 모델 저장위치 확인

훈련이 끝나면 위와 같은 메시지가 출력 되는데 runs/detect/train 위치에 last 모델과 best 모델이 저장되어 있다는 것을 확인할 수 있다.

4. 학습을 마친 후의 모델 정보 확인해 보자.

print(type(model.names),len(model.names))
print(model.names)

smoke 만 감지하는 모델로 변경이 된 것을 확인 할 수 있다.

 

6. 모델을 이용해서 예측하기

results=model(source='/content/dataset/test/images',save=True)

테스트 이미지를 모델을 이용해서 예측해 보자.

예측 결과는 train2 폴더에 저장 되어 있다.

데이터를 살펴 보면 위와 같이 산불의 연기를 감지하는 것을 알 수 있다.

 

7. 감지시 알람 메시지 출력하기

import os
import ultralytics
from ultralytics import YOLO

def check_fire_and_alert(results):
    for result in results:
      for box in result.boxes:
          class_id = box.cls.item()
          confidence = box.conf.item()
          # Check if the detected object is 'fire' (assuming 'fire' is at index 1 in your model.names)
          # Replace '1' with the actual index of the 'fire' class if different
          if confidence > 0.5: # Adjust confidence threshold as needed
              print("FIRE DETECTED!!!")
              print(f"Confidence: {confidence}")
              # Add further actions here, such as sending an alert message
              # Example: sending an email or SMS alert
              return # Exit the function after the first fire detection

# Example usage:
# Assuming 'results' is already populated from your model inference
check_fire_and_alert(results)

결과 값에 box가 있고 여기서 confidence 가 0.5 보다 크다면 FIRE DETECTED 를 출력하도록 하였다.

실제 애플리케이션이라면 CCTV 관제를 하다가 산불이 감지 되면 알람 발생으로 즉시 조치를 취할 수 있을 것이다.

 

소스코드 : https://colab.research.google.com/drive/1THTk7uH9gk87gmy7zOFU-4kR-jCHtWNp?usp=sharing