yolov8的推理和后处理代码

import onnxruntime as ort
import numpy as np
import cv2
import os

def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))

def softmax(x):
    x = x - np.max(x)
    e_x = np.exp(x)
    return e_x / e_x.sum()

def iou(a, b):
    cross_left = max(a[0], b[0])
    cross_top = max(a[1], b[1])
    cross_right = min(a[2], b[2])
    cross_bottom = min(a[3], b[3])
    cross_area = max(0.0, cross_right - cross_left) * max(0.0, cross_bottom - cross_top)
    union_area = max(0.0, a[2]-a[0])*max(0.0, a[3]-a[1]) + max(0.0, b[2]-b[0])*max(0.0, b[3]-b[1]) - cross_area
    if cross_area == 0 or union_area == 0:
        return 0.0
    return cross_area / union_area

def decode(score_ptr, reg_ptr, fea_w, fea_h, stride, class_num, reg_max, reg_scale_factor, conf_thresh, img_w, img_h):
    prebox = []
    dthreshold = -np.log(1.0 / conf_thresh - 1.0)
    for i in range(fea_h):
        for j in range(fea_w):
            scores = score_ptr[:class_num]
            score_ptr = score_ptr[class_num:]
            class_id = np.argmax(scores)
            max_conf = scores[class_id]
            if max_conf < dthreshold:
                reg_ptr = reg_ptr[4 * reg_max:]
                continue
            conf = sigmoid(max_conf)
            dxyxy = np.zeros(4, dtype=np.float32)
            for k in range(4):
                reg = reg_ptr[k*reg_max : (k+1)*reg_max]
                reg = softmax(reg)
                dxyxy[k] = (np.arange(reg_max) * reg).sum() * reg_scale_factor
            reg_ptr = reg_ptr[4 * reg_max:]
            x0 = (j + 0.5 - dxyxy[0]) * stride
            y0 = (i + 0.5 - dxyxy[1]) * stride
            x1 = (j + 0.5 + dxyxy[2]) * stride
            y1 = (i + 0.5 + dxyxy[3]) * stride
            x0 = np.clip(x0, 0, img_w)
            y0 = np.clip(y0, 0, img_h)
            x1 = np.clip(x1, 0, img_w)
            y1 = np.clip(y1, 0, img_h)
            prebox.append([x0, y0, x1, y1, class_id, conf])
    return prebox

def postprocess(img_path, yuv_name):
    YOLOV8_REG_MAX = 8
    YOLOV8_REG_MAX_EQ = 16
    class_num = 4
    conf_thresh = 0.3
    iou_thresh = 0.5
    img_w, img_h = 960, 512
    reg_scale_factor = YOLOV8_REG_MAX_EQ // YOLOV8_REG_MAX
    strides = [8, 16, 32]
    sizes = [(img_h//8, img_w//8), (img_h//16, img_w//16), (img_h//32, img_w//32)]
    
    cls_files = ["Cls_0.bin", "Cls_1.bin", "Cls_2.bin"]
    reg_files = ["Reg_0.bin", "Reg_1.bin", "Reg_2.bin"]
    
    all_bboxes = []
    for i in range(3):
        h, w = sizes[i]
        stride = strides[i]
        cls = np.fromfile(cls_files[i], dtype=np.float32)
        reg = np.fromfile(reg_files[i], dtype=np.float32)
        bboxes = decode(cls, reg, w, h, stride, class_num, YOLOV8_REG_MAX, reg_scale_factor, conf_thresh, img_w, img_h)
        all_bboxes.extend(bboxes)
    
    all_bboxes.sort(key=lambda x: x[5], reverse=True)
    n = len(all_bboxes)
    remove = [False]*n
    results = []
    
    for i in range(n):
        if remove[i]: continue
        box = all_bboxes[i]
        results.append(box)
        for j in range(i+1, n):
            if remove[j]: continue
            if box[4] == all_bboxes[j][4]:
                if iou(box, all_bboxes[j]) >= iou_thresh:
                    remove[j] = True

    print(yuv_name)
    print(len(results))
    for r in results:
        x0, y0, x1, y1, cid, prob = r
        x = int(x0)
        y = int(y0)
        w = int(x1 - x0)
        h = int(y1 - y0)
        print(f"DETECT: x={x}, y={y}, w={w}, h={h}, cid={int(cid)}, prob={prob:.4f}")

    img = cv2.imread(img_path)
    for r in results:
        x0, y0, x1, y1, cid, prob = r
        x0 = int(x0)
        y0 = int(y0)
        x1 = int(x1)
        y1 = int(y1)
        cv2.rectangle(img, (x0, y0), (x1, y1), (0, 255, 0), 2)
        # label = f"{int(cid)} {prob:.2f}"
        # cv2.putText(img, label, (x0, y0-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
    cv2.imwrite(img_path, img)

def process_single_yuv(yuv_path):
    h, w = 512, 960
    with open(yuv_path, "rb") as f:
        nv21_data = f.read()

    yuv = np.frombuffer(nv21_data, dtype=np.uint8).reshape(h * 3 // 2, w)
    rgb = cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB_NV21)
    
    base_name = os.path.splitext(yuv_path)[0]
    jpg_path = f"{base_name}.jpg"
    cv2.imwrite(jpg_path, cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR))

    input_tensor = rgb.astype(np.float32) / 255.0
    input_tensor = input_tensor.transpose(2, 0, 1)
    input_tensor = np.expand_dims(input_tensor, 0)

    session = ort.InferenceSession("240731_960_512_4_direct_split_r8_published_epoch_50_1332dcd5_float.onnx")
    input_name = session.get_inputs()[0].name
    output_names = [o.name for o in session.get_outputs()]

    outputs = session.run(output_names, {input_name: input_tensor})

    for name, out in zip(output_names, outputs):
        out.tofile(f"{name}.bin")

    postprocess(jpg_path, yuv_path)

def main():
    yuv_files = [
        "block_row0_col0.yuv", "block_row0_col1.yuv", "block_row0_col2.yuv", "block_row0_col3.yuv",
        "block_row1_col0.yuv", "block_row1_col1.yuv", "block_row1_col2.yuv", "block_row1_col3.yuv",
        "block_row2_col0.yuv", "block_row2_col1.yuv", "block_row2_col2.yuv", "block_row2_col3.yuv",
        "block_row3_col0.yuv", "block_row3_col1.yuv", "block_row3_col2.yuv", "block_row3_col3.yuv"
    ]
    
    for yuv in yuv_files:
        if os.path.exists(yuv):
            process_single_yuv(yuv)

if __name__ == "__main__":
    main()