6.ONNX 모델 검증

파이썬에서 위와 같이 입력해서 output 값을 확인하자.

import onnxruntime
import numpy as np

# ONNX 모델 로드
onnx_model_path = "C:/Users/L/runs/detect/train/weights/best.onnx"
session = onnxruntime.InferenceSession(onnx_model_path)

# 입력 이름 확인
input_name = session.get_inputs()[0].name
print(f"Input name: {input_name}")

# 출력 이름 확인
output_names = [output.name for output in session.get_outputs()]
print(f"Output names: {output_names}")

# 입력 데이터 생성 (더미 입력)
input_size = 640
dummy_input = np.random.rand(1, 3, input_size, input_size).astype(np.float32)  # 랜덤 입력

# 모델 추론
outputs = session.run(None, {input_name: dummy_input})

# 출력 구조 확인
for idx, output in enumerate(outputs):
    print(f"Output {idx}: shape = {np.array(output).shape}")

위에서 나온결과는 Output 0: shape = (1, 5, 8400) 인데

첫번째1 은 입력이미지 크기이고 일반적으로 1이다. 이건 넘어가고

두번째 숫자 5 와 8400이 중요 하다.

5는 바운딩박스 정보 , 즉 감지한 박스의 정보인데 x, y, w, h, confidence 순으로 들어있다는 의미이다.

모델의 감지 결과는 결과는 한줄로 늘어져있는 1차원 배열 형태로 나오는데 거기에 저값들이 순차적으로 들어있다.

그런데 반복해서 들어있는 형태가 특이한데 5다음 숫자가 8400이므로 [x1,x2,x3,x4....x8400,y1,y2,y3,y4...y8400,w1,w2,w3....w8400,h1,h2,h3....h8400,c1,c2......c8400] 이런형태로 들어있다.

즉 감지된 첫번째 박스를 확인하려면 결과값의 첫번째값 = x ,8400번째값 = y , 16800번째값 = w , 25200번째값 = h , 33600 번째값이 confidence(딥러닝 스코어 0~1) 이 된다.

배열을 5개로 나눠서 첫번째 값들끼리 합쳐야 첫번째 결과, 두번째 값들끼리 합쳐야 두번째 결과가 된다는 말이다.

위 값들을 기억해 두고 이제 테스트를 해보자.

 

 

 

위와 같이 pt 파일이 정상적으로 대상을 잡는지 확인부터 해보자.

100장정도로도 상당히 좋은 스코어가 나오는것을 볼수 있다.

 

 

from ultralytics import YOLO
import cv2
import matplotlib.pyplot as plt

# 모델 로드 
model_path = "C:/Users/L/runs/detect/train/weights/best.pt"
model = YOLO(model_path)  # 모델 로드

# 이미지 로드
image_path = "C:/Users/L/Desktop/dog/2.png"
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

# 모델을 이용한 추론
results = model(image_rgb, conf=0.9)  # 신뢰도 임계값이 낮을수록 많은탐지가능

# 결과 확인
for result in results:
    # 바운딩 박스 확인
    print("Boxes:", result.boxes)  # 바운딩 박스 정보 출력
    
    if result.boxes is not None and len(result.boxes.xyxy) > 0:
        print("Boxes (xyxy):", result.boxes.xyxy[0])  # x1, y1, x2, y2
        print("Scores:", result.boxes.conf[0])  # 신뢰도 점수
        print("Class IDs:", result.boxes.cls[0])  # 클래스 ID
    else:
        print("No detections.")
    
    # 시각화: 바운딩 박스 시각화
    annotated_img = result.plot()  # 원본 이미지에 바운딩 박스 시각화
    plt.imshow(annotated_img)
    plt.axis('off')
    plt.show()

위에서 테스트한 runpt.py 파일 소스이다.

 

 

 

 

위와같이 변경한 onnx 파일로도 대상을 정확하게 잡는지 확인해보자.

 

cmd창에서

pip install opencv-python numpy onnxruntime matplotlib

입력해서 설치 한다.

 

 

import cv2
import numpy as np
import onnxruntime as ort
import matplotlib.pyplot as plt

# ONNX 모델 경로
onnx_model_path = "C:/Users/L/runs/detect/train/weights/best.onnx"

# ONNX 세션 생성
session = ort.InferenceSession(onnx_model_path)

# 이미지 로드
image_path = "C:/Users/L/Desktop/dog/2.png"
image = cv2.imread(image_path)
orig_h, orig_w = image.shape[:2]

# 1. 전처리
input_size = 640
image_resized = cv2.resize(image, (input_size, input_size))
image_rgb = cv2.cvtColor(image_resized, cv2.COLOR_BGR2RGB)
input_data = image_rgb.astype(np.float32) / 255.0
input_data = np.transpose(input_data, (2, 0, 1))
input_data = np.expand_dims(input_data, axis=0)

# 2. 모델 추론
input_name = session.get_inputs()[0].name
outputs = session.run(None, {input_name: input_data})
output = outputs[0]  # (1, 5, 8400)

# 출력 디버깅
print("Output shape:", output.shape)
print("Sample output:", output[0, :, :5])

# 3. 출력 후처리
boxes = output[0, :4, :]  # 바운딩 박스 좌표 (cx, cy, w, h)
confidences = output[0, 4, :]  # 객체 신뢰도

# 좌표 변환 및 필터링
scaled_boxes = []
final_scores = []
for i in range(boxes.shape[1]):
    cx, cy, w, h = boxes[:, i]
    if confidences[i] > 0.6:  # 신뢰도 임계값 적용
        # 중심 좌표 (cx, cy) 및 너비/높이 (w, h) 변환
        x1 = int(cx - w / 2)  # 좌상단 x
        y1 = int(cy - h / 2)  # 좌상단 y
        x2 = int(cx + w / 2)  # 우하단 x
        y2 = int(cy + h / 2)  # 우하단 y

        # 스케일 조정: 모델 입력 크기 -> 원본 이미지 크기
        x1 = int(x1 * orig_w / input_size)
        y1 = int(y1 * orig_h / input_size)
        x2 = int(x2 * orig_w / input_size)
        y2 = int(y2 * orig_h / input_size)

        scaled_boxes.append([x1, y1, x2, y2])
        final_scores.append(confidences[i])

# NMS(비최대 억제)
indices = cv2.dnn.NMSBoxes(scaled_boxes, final_scores, 0.6, 0.4)

# 4. 결과 시각화
if len(indices) > 0:
    for i in indices.flatten():
        x1, y1, x2, y2 = scaled_boxes[i]
        label = f"Conf: {final_scores[i]:.2f}"
        cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)

# 결과 표시
plt.figure(figsize=(10, 8))
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
plt.axis('off')
plt.show()

위에서 테스트한 run.py 파일이다.