Ultralytics의 YOLOv8 프레임워크는 분류, 탐지, 분할 등 다양한 비전 작업을 한 줄의 코드로 학습하고 배포할 수 있는 강력한 기능을 제공합니다. 하지만 .pt 가중치를 ONNX 형식으로 내보낸 후, ultralytics 라이브러리가 없는 환경에서 onnxruntime만으로 추론을 시도하면 결과값이 일치하지 않거나 신뢰도(Confidence)가 현저히 낮게 측정되는 경우가 발생합니다.

이는 학습 시 적용된 특정 전처리(Preprocessing) 로직이 ONNX 모델에 포함되지 않기 때문입니다. 따라서 일관된 추론 결과를 얻으려면 YOLOv8 분류 작업의 내부 전처리 로직을 직접 구현해야 합니다.

YOLOv8 분류 전처리 로직의 이해

YOLOv8 분류 모델의 전처리는 입력 이미지의 크기 조정, 중앙 크롭(Center Crop), 그리고 정규화 과정을 포함합니다. 특히 crop_fraction 값을 고려하여 이미지를 먼저 확장된 크기로 리사이즈한 뒤 중앙을 잘라내는 과정이 핵심입니다. 아래는 이를 PyTorch의 torchvision.transforms를 사용하여 재현한 코드입니다.

import torch
import math
from typing import Tuple, Union
import torchvision.transforms as transforms
from torchvision.transforms import InterpolationMode

def build_yolo_preprocessor(
    input_size: int = 224, 
    mean: Tuple[float, ...] = (0.0, 0.0, 0.0), 
    std: Tuple[float, ...] = (1.0, 1.0, 1.0), 
    crop_ratio: float = 1.0
) -> transforms.Compose:
    """
    YOLOv8 분류 모델의 기본 전처리 과정을 재현합니다.
    """
    # 리사이즈 대상 크기 계산
    if isinstance(input_size, (list, tuple)):
        target_size = tuple(math.floor(x / crop_ratio) for x in input_size)
    else:
        dimension = math.floor(input_size / crop_ratio)
        target_size = (dimension, dimension)

    pipeline = []
    
    # 1. 리사이즈 단계
    if target_size[0] == target_size[1]:
        pipeline.append(transforms.Resize(target_size[0], interpolation=InterpolationMode.BILINEAR))
    else:
        pipeline.append(transforms.Resize(target_size, interpolation=InterpolationMode.BILINEAR))

    # 2. 중앙 크롭 및 텐서 변환, 정규화
    pipeline.extend([
        transforms.CenterCrop(input_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=torch.tensor(mean), std=torch.tensor(std))
    ])

    return transforms.Compose(pipeline)

전체 워크플로우: 학습부터 ONNX 추론까지

1. 모델 학습 및 ONNX 내보내기

먼저 모델을 학습시키고 배포용 ONNX 파일로 변환합니다.

from ultralytics import YOLO

# 모델 초기화 및 학습
model = YOLO("yolov8n-cls.pt")
model.train(data="custom_dataset_path", epochs=50, imgsz=224)

# 최적 가중치를 ONNX로 변환
best_model = YOLO("runs/classify/train/weights/best.pt")
best_model.export(format="onnx")

2. ONNX Runtime을 이용한 추론 구현

이제 ultralytics 의존성 없이 전처리 로직을 결합하여 추론을 수행하는 전체 코드입니다.

import cv2
import numpy as np
import onnxruntime as ort
from PIL import Image
import torch
import torchvision.transforms as transforms

def execute_inference():
    # 설정 값
    MODEL_PATH = "best.onnx"
    IMAGE_PATH = "test_image.jpg"
    CLASS_NAMES = {0: "category_A", 1: "category_B"}
    
    # 세션 로드 (GPU 사용 가능 시 CUDA 사용)
    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    session = ort.InferenceSession(MODEL_PATH, providers=providers)
    
    # 전처리 파이프라인 생성
    preprocess = build_yolo_preprocessor(input_size=224)

    # 이미지 로드 및 전처리 적용
    raw_img = cv2.imread(IMAGE_PATH)
    rgb_img = cv2.cvtColor(raw_img, cv2.COLOR_BGR2RGB)
    pil_img = Image.fromarray(rgb_img)
    
    # 배치 차원 추가 및 Numpy 변환
    input_tensor = preprocess(pil_img).unsqueeze(0).numpy()

    # ONNX 모델 입출력 이름 획득
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    # 추론 실행
    results = session.run([output_name], {input_name: input_tensor})[0]
    
    # 결과 해석
    top_idx = int(np.argmax(results[0]))
    top_score = float(results[0][top_idx])

    print(f"Result: {CLASS_NAMES.get(top_idx, top_idx)} ({top_score:.4f})")

if __name__ == "__main__":
    execute_inference()

이와 같이 수동으로 전처리 파이프라인을 구축하면, ultralytics 라이브러리를 설치할 수 없는 임베디드 환경이나 경량화된 프로덕션 서버에서도 YOLOv8 분류 모델을 정확하고 안정적으로 운영할 수 있습니다.