from ultralytics import YOLO
from ultralytics.utils.ops import non_max_suppression
import torch
import cv2

# ======================
# 配置参数
# ======================
MODEL_PATH = '/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_ailai_detect2/weights/best.pt'
IMG_PATH = '3.jpg'
OUTPUT_PATH = 'output_pt.jpg'
CONF_THRESH = 0.5
IOU_THRESH = 0.45
CLASS_NAMES = ['bag', 'bag35']

# ======================
# 主函数
# ======================
def main():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    print(f"✅ 使用设备: {device}")

    model = YOLO(MODEL_PATH).to(device)

    print("➡️ 开始推理...")
    results = model(IMG_PATH, imgsz=640, conf=CONF_THRESH, device=device, verbose=True)

    r = results[0]
    pred = r.boxes.data  # GPU tensor [N,6]

    det = non_max_suppression(
        pred.unsqueeze(0),
        conf_thres=CONF_THRESH,
        iou_thres=IOU_THRESH,
        classes=None,
        agnostic=False,
        max_det=100
    )[0]

    if det is None or len(det) == 0:
        print("❌ 未检测到任何目标")
        return

    det = det.cpu().numpy()  # 只拷贝一次

    # ======================
    # ⭐ 关键：取置信度最高的结果
    # ======================
    best_det = max(det, key=lambda x: x[4])

    x1, y1, x2, y2, conf, cls_id = best_det
    x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
    cls_id = int(cls_id)
    cls_name = CLASS_NAMES[cls_id]

    print("\n🏆 置信度最高结果:")
    print(f"  类别: {cls_name}")
    print(f"  置信度: {conf:.3f}")
    print(f"  框: [{x1}, {y1}, {x2}, {y2}]")

    # ======================
    # 可视化（只画最高的）
    # ======================
    img = cv2.imread(IMG_PATH)
    if img is None:
        raise FileNotFoundError(f"无法读取图像: {IMG_PATH}")

    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
    label = f"{cls_name} {conf:.2f}"
    cv2.putText(
        img,
        label,
        (x1, max(y1 - 10, 0)),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.9,
        (0, 255, 0),
        2
    )

    cv2.imwrite(OUTPUT_PATH, img)
    print(f"\n🖼️ 可视化结果已保存: {OUTPUT_PATH}")

if __name__ == '__main__':
    main()