dsh/testrk3588/test_onnx.py

#!/usr/bin/env python3
import cv2
import numpy as np
import onnxruntime as ort

LABEL_MAP = {0: '安全帽', 4: '安全衣', 3: '人'}

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

def test_onnx():
    model_path = "yolo11m_safety.onnx"
    input_size = (640, 640)
    
    image = cv2.imread("b.jpg")
    if image is None:
        print("无法读取测试图片")
        return
    
    h0, w0 = image.shape[:2]
    
    print(f"=== 预处理 ===")
    img = cv2.resize(image, input_size)
    img = img[..., ::-1].astype(np.float32) / 255.0
    img = img.transpose(2, 0, 1)
    blob = img[None, ...]
    print(f"Input blob shape: {blob.shape}")
    
    print(f"\n=== 加载 ONNX 模型 ===")
    session = ort.InferenceSession(model_path)
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name
    print(f"Input name: {input_name}")
    print(f"Output name: {output_name}")
    
    print(f"\n=== 推理 ===")
    outputs = session.run([output_name], {input_name: blob})
    output = outputs[0]
    print(f"Output shape: {output.shape}")
    
    print(f"\n=== 原始输出分析 ===")
    output_a = output[0]
    print(f"After squeeze shape: {output_a.shape}")
    
    print(f"\n=== 查找高置信度框 (obj_conf > 0.1) ===")
    high_obj_indices = []
    for i in range(output_a.shape[1]):
        obj_conf = output_a[4, i]
        if obj_conf > 0.1:
            high_obj_indices.append((i, obj_conf))
    
    high_obj_indices.sort(key=lambda x: x[1], reverse=True)
    print(f"找到 {len(high_obj_indices)} 个高置信度框")
    
    print(f"\n前 20 个高置信度框:")
    for idx, obj_conf in high_obj_indices[:20]:
        x_center = float(output_a[0, idx])
        y_center = float(output_a[1, idx])
        width = float(output_a[2, idx])
        height = float(output_a[3, idx])
        class_probs_raw = output_a[5:9, idx]
        class_probs = sigmoid(class_probs_raw)
        class_id = int(np.argmax(class_probs))
        class_conf = float(class_probs[class_id])
        confidence = obj_conf * class_conf
        
        print(f"\nBox {idx}:")
        print(f"  坐标: x={x_center:.3f}, y={y_center:.3f}, w={width:.3f}, h={height:.3f}")
        print(f"  置信度: obj_conf={obj_conf:.3f}, class_conf={class_conf:.6f}, total={confidence:.3f}")
        print(f"  类别: class_id={class_id}, name={LABEL_MAP.get(class_id, 'unknown')}")
    
    print(f"\n=== 检测结果 ===")
    dets = []
    h0, w0 = image.shape[:2]
    for idx, obj_conf in high_obj_indices:
        x_center = float(output_a[0, idx])
        y_center = float(output_a[1, idx])
        width = float(output_a[2, idx])
        height = float(output_a[3, idx])
        class_probs_raw = output_a[5:9, idx]
        class_probs = sigmoid(class_probs_raw)
        class_id = int(np.argmax(class_probs))
        class_conf = float(class_probs[class_id])
        confidence = obj_conf * class_conf
        
        if class_id not in LABEL_MAP:
            continue
        
        conf_threshold = 0.01
        if confidence < conf_threshold:
            continue
        
        x1 = int(x_center - width/2)
        y1 = int(y_center - height/2)
        x2 = int(x_center + width/2)
        y2 = int(y_center + height/2)
        
        x1 = int(x1 * (w0/640))
        y1 = int(y1 * (h0/640))
        x2 = int(x2 * (w0/640))
        y2 = int(y2 * (h0/640))
        
        x1 = max(0, x1)
        y1 = max(0, y1)
        x2 = min(w0, x2)
        y2 = min(h0, y2)
        
        dets.append({
            'class_id': class_id,
            'class_name': LABEL_MAP[class_id],
            'confidence': confidence,
            'bbox': (x1, y1, x2, y2)
        })
    
    print(f"\n检测到 {len(dets)} 个框 (阈值 0.01)")
    for d in dets[:20]:
        print(f"  {d['class_name']}: conf={d['confidence']:.3f}, box={d['bbox']}")

if __name__ == "__main__":
    test_onnx()