#!/usr/bin/env python3
import cv2
import numpy as np
from rtsp_person_detection import create_detector

# 测试检测功能
model_path = 'yolo11m_safety.onnx'  # 可以切换为 'yolo11m_safety.rknn'
det = create_detector(model_path)

# 读取测试图片
img = cv2.imread('b.jpg')
if img is None:
    print('无法读取测试图片')
    exit(1)

print(f'测试图片形状: {img.shape}')

# 进行检测
conf_threshold_map = {3: 0.8, 0: 0.5, 4: 0.5}
detections = det.detect(img, conf_threshold_map)

print(f'检测结果: {len(detections)} 个目标')
for det in detections:
    print(f'  {det.class_name}: conf={det.confidence:.2f}, box={det.bbox}')