Yolo를 onnx로 변환하여 Unity에 심어 보려고 테스트를 하던도중 일부 시행착오가 있어 공유하고자 기재하여본다. Yolo를 사용하여 자신에 데이터로 학습하고 onnx파일로 만드는 것 부터 하고 싶은 분은 아래 링크를 참고하기 바란다.

Yolo v8 사용예제
Model Load from ultralytics import YOLOimport matplotlib.pyplot as pltimport cv2,torchmodel = YOLO(“yolov8n.yaml”) # build a new model from scratchmodel = YOLO(“yolov8n.pt”) # load a pretrained model (recommended for training) # pre training model.train(data=“data.yaml”,epochs=1,batch=2)ne…
IT 기술블로그 💡조창제

ONNX의 Inference 값과 pytorch 모델의 Inference 값이 달라 무엇이 문제인가 검색해보던 도중 아래 블로그에서 디텍션 부분까지 잘 설명이 되어있었다. 하지만 하나의 물체에 여러 박스가 겹치는 문제인 Multiple detections에 대한 개선방법인 NMS(Non Maximal Suppression)이 기재되어 있지 않아 같이 포스팅해보았다.

How to load model YOLOv8 ONNXRuntime
YOLOv8 is the latest version (v8) of the YOLO (You Only Look Once) object detection system. YOLO is a real-time, one-shot object detection…
MediumAli Mustofa

NMS는 아래에서 잘 설명되고 있는데 그 중 나는 여러 오브젝트가 중복될 일이 없기 때문에 Yolo에서 사용되고 있다는 NMS방식을 구현하였다.

3) NMS (Non-Maximum Suppression) & Anchor box
## 서론 앞서 [1) General process of object detection](https://wikidocs.net/152376) 에서 언급했듯, 대다수의 objec…
위키독스

Package Load & Model Load

import matplotlib.pyplot as plt
import numpy as np
import onnx,onnxruntime,time,tqdm,os,cv2
from my_utils import * # my_utils는 맨 아래에 있음

onnx_model = onnx.load("my_weight.onnx")
onnx.checker.check_model(onnx_model)

opt_session = onnxruntime.SessionOptions()
opt_session.enable_mem_pattern = False
opt_session.enable_cpu_mem_arena = False
opt_session.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL

ort_session = onnxruntime.InferenceSession(
  "my_weight.onnx",providers=[
    'CUDAExecutionProvider', 'CPUExecutionProvider'])

Model Size Check

model_inputs = ort_session.get_inputs()
input_names = [model_inputs[i].name for i in range(len(model_inputs))]
input_shape = model_inputs[0].shape # [1, 3, 640, 640] yolo model input size
model_output = ort_session.get_outputs()
output_names = [model_output[i].name for i in range(len(model_output))]

Pre Processing & Prediction

image=cv2.imread("WIN_20230424_14_20_31_Pro.jpg") #file read
image_height, image_width = image.shape[:2]
input_height, input_width = input_shape[2:] # 640, 640
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) 
resized = cv2.resize(image_rgb, (input_width, input_height)) # 640,640

# Scale input pixel value to 0 to 1
input_image = resized / 255.0
input_image = input_image.transpose(2,0,1)
input_tensor = input_image[np.newaxis, :, :, :].astype(np.float32)
input_tensor.shape

outputs = ort_session.run(output_names, {input_names[0]: input_tensor})[0]
predictions = np.squeeze(outputs).T

Post Processing

conf_thresold = 0.25 # Default
# Filter out object confidence scores below threshold
scores = np.max(predictions[:, 4:], axis=1)
predictions = predictions[scores > conf_thresold, :]
scores = scores[scores > conf_thresold]  

# NMS algorithm
sorted_pred=predictions[np.argsort(predictions[:,4])]
sorted_pred[:,:4]=xywh2xyxy(sorted_pred[:,:4])

view=image[:,:,::-1].copy()
for bbox in boxes:
    bbox = bbox.round().astype(np.int32).tolist()
    plt.imshow(
      cv2.rectangle(
        view, tuple(bbox[:2]), tuple(bbox[2:]), (255,0,0), 2))

utils

# My utils
import numpy as np 
def xywh2xyxy(x):
    # Convert bounding box (x, y, w, h) to bounding box (x1, y1, x2, y2)
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2
    y[..., 1] = x[..., 1] - x[..., 3] / 2
    y[..., 2] = x[..., 0] + x[..., 2] / 2
    y[..., 3] = x[..., 1] + x[..., 3] / 2
    return y

def IoU(box1, box2):
    # box = (x1, y1, x2, y2)
    box1_area = (box1[2] - box1[0] + 1) * (box1[3] - box1[1] + 1)
    box2_area = (box2[2] - box2[0] + 1) * (box2[3] - box2[1] + 1)

    # obtain x1, y1, x2, y2 of the intersection
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])

    # compute the width and height of the intersection
    w = max(0, x2 - x1 + 1)
    h = max(0, y2 - y1 + 1)

    inter = w * h
    iou = inter / (box1_area + box2_area - inter)
    return iou


def check_iou(pred):
    base=pred[0]
    keep_list=list()
    for i in range(1,pred.shape[0]):
        iou_score=IoU(base, pred[i])
        if iou_score<=0.5:
            keep_list.append(i)
    return keep_list

def NMS(pred):
    """
    NMS : Non Maximum Suppression
    """
    out=list()
    while pred.shape[0]!=0:
        out.append(pred[0])
        keep_idx=check_iou(pred[:,:4])
        pred=pred[keep_idx]
    return np.array(out)