[Python] YOLOv8.pt를 ONNX 모델로 export, inference

YOLOv8n 모델을 onnx모델로 export합니다.

from ultralytics import YOLO
model = YOLO("yolov8n.pt")
model.export(format="onnx")

yolov8n.onnx 파일이 생성됩니다.

ONNX 모델로 Detection

라이브러리 설치후 import합니다.

!pip install onnxruntime
import onnxruntime as ort

모델을 불러옵니다.

model = ort.InferenceSession("yolov8n.onnx")

모델 RUN

outputs = model.get_outputs()
output = outputs[0]
outputs = model.run(["output0"], {"images":input})

output_names - 수신하고자 하는 출력의 이름 배열입니다. YOLOv8 모델에서는 단일 항목이 있는 배열이 됩니다.
input - input의 dictionary로, {name:tensor} 형식으로 네트워크에 전달합니다. 여기서 name은 input의 이름이고 tensor는 앞서 준비한 이미지 데이터 배열입니다.

만약 input 형식 에러가 뜨면 형변환을 합니다.
input = input.astype(np.float32)

output에서 좌표를 추출합니다.

output = outputs[0]
output = output[0]
output = output.transpose()
row = output[0]
xc,yc,w,h = row[:4]

row는[x1,y1,x2,y2,class_label,probability] 형식으로 구성되어 있습니다.

x1 = xc-w/2
y1 = yc-h/2
x2 = xc+w/2
y2 = yc+h/2

x1 = (xc - w/2) / 640 * img_width
y1 = (yc - h/2) / 640 * img_height
x2 = (xc + w/2) / 640 * img_width
y2 = (yc + h/2) / 640 * img_height

prob = row[4:].max()
class_id = row[4:].argmax()

추출한 좌표로 박스를 그립니다.

frame = cv2.imread("bus.jpg")

for box in boxes:
    x1,y1,x2,y2,class_id,prob = box
    cv2.rectangle(frame, pt1= (int(x1),int(y1)),pt2=(int(x2),int(y2)),color=(0, 0, 255), thickness=2)