OpenCV 강좌, ONNX 형식으로 모델 변환 후 OpenCV에서 실행

1. ONNX란 무엇인가?

ONNX는 다양한 머신러닝 프레임워크 간의 상호 운용성을 높이기 위해 설계된 개방형 포맷입니다. ONNX 형식을 사용하면 TensorFlow, PyTorch 등 여러 프레임워크에서 학습한 모델을 손쉽게 변환하고 사용할 수 있습니다. ONNX는 다음과 같은 장점을 제공합니다:

• 모델 호환성: 다양한 플랫폼과 언어에서 쉽게 로드하는 것이 가능하다.
• 퍼포먼스: 최적화된 연산을 통해 모델의 실행속도를 향상시킬 수 있다.
• 유연성: 다중 프레임워크 지원으로 인해 특정 프레임워크에 국한되지 않는다.

2. OpenCV란 무엇인가?

OpenCV(Open Source Computer Vision Library)는 컴퓨터 비전 및 머신러닝을 위한 라이브러리로서, C++, Python, Java 등 다양한 언어를 지원합니다. OpenCV는 이미지 처리, 비디오 분석, 비디오 감시, 로봇 비전 등의 분야에서 광범위하게 사용됩니다.

3. ONNX 모델 변환 과정

3.1. PyTorch 모델을 ONNX로 변환하기

PyTorch에서 훈련한 모델을 ONNX 형식으로 변환하는 방법은 다음과 같습니다. 예를 들어, 간단한 이미지 분류 모델을 만들어 보겠습니다.

import torch
import torch.onnx
import torchvision.models as models

# 기본 ResNet 모델을 생성
model = models.resnet18(pretrained=True)
model.eval()

# 더미 입력 데이터 생성 (배치 크기 1, 채널 3, 높이 224, 너비 224)
dummy_input = torch.randn(1, 3, 224, 224)

# ONNX 형식으로 모델 저장
torch.onnx.export(model, dummy_input, "resnet18.onnx", 
                  verbose=True, 
                  input_names=['input'], 
                  output_names=['output'])
        

위 코드에서는 PyTorch에서 제공하는 미리 훈련된 ResNet-18 모델을 사용하여 .onnx 파일로 변환하고 있습니다. torch.onnx.export 함수의 매개변수로 모델, 더미 입력 데이터, 저장할 파일 이름을 지정합니다. 또한 입력 및 출력 텐서의 이름을 정의할 수 있습니다.

3.2. TensorFlow 모델을 ONNX로 변환하기

TensorFlow에서 훈련된 모델도 ONNX 형식으로 변환이 가능합니다. TensorFlow에서 ONNX로 변환하려면 tf2onnx 라이브러리를 사용할 수 있습니다. 다음은 TensorFlow 모델을 ONNX 형식으로 변환하는 예제입니다.

import tensorflow as tf
import tf2onnx

# 예제 모델 생성
model = tf.keras.applications.MobileNetV2(weights='imagenet', input_shape=(224, 224, 3))

# 입력 및 출력 이름 정의
output_path = "mobilenetv2.onnx"
model_tensor = tf.convert_to_tensor(tf.ones([1, 224, 224, 3]), dtype=tf.float32)

# ONNX 모델로 변환
onnx_model = tf2onnx.convert.from_keras(model)
tf2onnx.save_model(onnx_model, output_path)
        

이 코드는 Keras를 사용하여 MobileNetV2 모델을 생성하고 이를 ONNX 파일로 변환합니다. tf2onnx.convert.from_keras() 메소드를 사용하여 모델을 ONNX로 변환하고, tf2onnx.save_model() 메소드로 저장합니다.

4. OpenCV에서 ONNX 모델 실행하기

ONNX 파일이 생성되면, OpenCV를 사용하여 이 모델을 실행할 수 있습니다. OpenCV는 ONNX 모델을 로드하고 입력 이미지를 모델에 전달하여 예측 결과를 얻는 기능을 제공합니다.

4.1. OpenCV 설치

OpenCV를 설치하려면 pip를 사용하여 설치할 수 있습니다. 아래 명령어를 통해 OpenCV를 설치합니다.

pip install opencv-python
pip install opencv-python-headless
        

4.2. ONNX 모델을 OpenCV로 로드하고 실행하기

아래 예제는 OpenCV에서 ONNX 모델을 로드하여 이미지를 분류하는 방법을 보여줍니다. 이 예제는 이전에 생성한 resnet18.onnx 파일을 사용합니다.

import cv2
import numpy as np

# ONNX 모델 로드
net = cv2.dnn.readNetFromONNX('resnet18.onnx')

# 입력 이미지 로드 및 전처리
input_image = cv2.imread('input_image.jpg')
input_image = cv2.resize(input_image, (224, 224))
blob = cv2.dnn.blobFromImage(input_image, 1.0, (224, 224), (104, 117, 123), swapRB=True, crop=False)

# 네트워크에 입력 데이터 설정
net.setInput(blob)

# 모델 실행 및 결과 얻기
output = net.forward()

# 결과 후처리 (가장 높은 확률의 클래스 찾기)
class_id = np.argmax(output[0])
confidence = output[0][class_id]

print(f"Class ID: {class_id}, Confidence: {confidence:.2f}")
        

위 코드에서는 OpenCV의 dnn 모듈을 사용하여 ONNX 모델을 로드하고, 입력 이미지를 전처리한 후 모델로 예측을 수행합니다. cv2.resize() 메소드로 이미지를 224×224 크기로 조정하고, cv2.dnn.blobFromImage() 메소드로 이미지를 블롭 형식으로 변환합니다. 마지막으로, net.forward() 메소드를 호출하여 모델의 예측 결과를 얻습니다.

5. 결론

ONNX 형식으로의 모델 변환과 이 모델을 OpenCV에서 실행하는 과정은 간단하지만 매우 강력한 기능을 제공합니다. 이를 통해 우리는 다양한 머신러닝 프레임워크에서 훈련된 모델을 손쉽게 활용할 수 있으며, OpenCV의 다양한 이미지 처리 기능을 모델 결과에 적용할 수 있습니다. 본 강좌를 통해 여러분이 OpenCV와 ONNX를 함께 사용하는 방법을 이해하고, 보다 효율적인 데이터 처리를 할 수 있기를 바랍니다.