Yolo모델을 이용한 오브젝트 분류(Object Classification) - 학습 및 실행

 

이전글에서 Yolo모델을 통해서 오브젝트 인식을 해보았습니다.  

Yolo모델을 이용한 오브젝트 인식(Object Detection) - 학습 및 실행

 

이번에 해볼 것은 오브젝트 분류!

 

인식과.. 분류..?

차이가 뭐지....?

 

 

인식은 어떠한 이미지나 환경에서 특정 오브젝트가 어디에 어떻게 있는지를 찾는것!

분류는 들어온 이미지가 어떤 오브젝트인지 선별하는것 입니다..

 

더더 단순히 표현하면,

오브젝트 분류는 이 오브젝트가 뭐지?,

오브젝트 인식은 어떠한 위치에 무언가 있네?(Localization)  = >>  이 오브젝트가 뭐지?(Classification)

 

이런 느낌인 것 같습니다...

 

 

Object Detection에 대한 그림과 자세한 설명은 아래 참고 블로그!

[Object Detection] 개념 정리

 

 

그럼 여기서 의문이 생깁니다...

엥 그러면 분류를 왜 쓰지, 인식으로 끝나는거 아닌가??

 => 그렇죠? 하지만 인식모델을 사용할때보다, 분류모델만 사용하는게 더더 빠르게 결과값을 받을수 있다는!

 

그래서! 이번엔 저번에 사용했던 Yolo모델을 통해 오브젝트 분류를 한번 해보려고 합니다!

 

 

 

그럼 시작!

---

 

우선 Yolo 모델 부터 확인해봅시다.

YOLO11 🚀 NEW

 

설명은 빨리빨리 넘어가고....

 

 

이번에 사용해볼 Classification 모델!

 

음... 구성은 동일 합니다...

 

Detection모델과 비교를 좀 해보면 

ImageSize가 224로 변했고....

속도가;; 엄청 빨라지네요....

이미지가 작아서 그런것도 있지만, 위에서 말했듯이 이미지 분류에만 모델을 사용하기에 빠를수 밖에 없는것 같습니다. 

 

 

---

그럼 다음 데이터!

 

Classification모델은... 따로 Yolo 형식의 txt가 필요하지 않습니다!

그리고 데이터의 위치, 정보, 이름, 클래스 등등이 포함되던 yaml파일도 필요가 없는!!

 

대신.. 아래처럼 디렉토리를 만들어 주셔야합니다.

ClassDataSet/
│── train/
│   ├── Apple/        
│   ├── Kiwi/         
│   ├── Watermelon/  
│── test/
│   ├── Apple/       
│   ├── Kiwi/         
│   ├── Watermelon/

 

 

 

그래서 기존에 라벨링 해주었던 데이터의 클래스들을 224*224에 맞게 잘라준뒤

 

 

 

아래와 같은 디렉토리 구조로 만들어 주었습니다. 

 

 

 

그럼이제 준비가 끝났습니다!

 

 

학습하러 ㄱㄱ!

 

---

학습코드!

코드는 아래와 같이... 

import os
import random
import numpy as np
from ultralytics import YOLO

# Step 1: 랜덤 시드 고정 (재현성 확보)
random.seed(821)
np.random.seed(821)

# Step 2: 데이터셋 폴더 경로 설정
dataset_path = os.path.dirname(os.path.realpath(__file__))

train_path = os.path.join(dataset_path, "Custom_fruits")  # 훈련 데이터 경로


# Step 3: YOLO Classification 모델 로드
model = YOLO("yolo11m-cls.pt")  # YOLOv8 분류 모델

# Step 4: 학습 시작
results = model.train(
    data=train_path,  	# 학습 데이터 경로
    epochs=100,  		# 학습 반복 횟수
    imgsz=224,  		# 이미지 크기
    batch=16,  			# 배치 크기
    workers=2,  		# 데이터 로딩 스레드 개수
    device=2,  			# GPU ID
    project=f"{dataset_path}/runs_cls", # 결과 저장 경로
    name="fruit_classifier",  			# 실험 이름
    save=True,  						# 학습 결과 저장 여부
    optimizer="AdamW",  				# AdamW 최적화 적용
    verbose=True  						# 학습 진행 로그 출력
)

 

경로를 설정해주고, 모델을 받아주고, model.train()함수로 끝!

 

 

굳... 학습이 잘 진행되는중

 

 

 

---

학습 끝! 오브젝트 분류 확인!

 

학습이 끝나면 위에서 지정해둔 경로상의 weight위치에 best.pt와 Last.pt의 모델이 생성됩니다.

이 모델을 이용해서 아래와 같은 코드를 작성하면!

import os
import torch
from ultralytics import YOLO
import cv2
import numpy as np
import pandas as pd

pathDir = "~~~~~~~" # 분류할 이미지들의 디렉토리경로
model_path = "~~~~~.pt"  # 학습된 모델 경로
model = YOLO(model_path)


DetectingList = ["Apple", "BlueBerry", "GreenApple", "Kiwi", "Lemon",  "Melon", "Orange", "Peach", "Pear", "Punica", "Watermelon"]

SaveData={}
for name in DetectingList:
    SaveData[name]= {"Ori_T": 0, "Ori_F": 0,"Pro_T": 0, "Pro_F": 0, "Gray_T": 0, "Gray_F": 0, "Gen_T": 0, "Gen_F": 0}

def YoloClassfication(image_paths):
    for img_path in image_paths:
        img = cv2.imread(img_path)

        # 모델 예측
        results = model(img)

        # 예측 결과 출력
        pred_label = results[0].probs.top1  # 가장 확률 높은 클래스 ID
        pred_confidence = results[0].probs.top1conf  # 해당 클래스의 확률
        class_name = model.names[pred_label]  # 클래스 이름


        print(f"이미지: {os.path.basename(img_path)}  예측: {class_name}")
        
        
  if __name__ == "__main__":
  		YoloClassfication(pathDir)

 

 

학습되어진 모델을 통해 분류결과를 확인 할 수 있습니다!

 

 

그럼 오브젝트 인식에서 분류까지 Yolo 실습 끝!

 

 

틀린점이 있다면 댓 달아주세요!