85 lines
2.2 KiB
Python
85 lines
2.2 KiB
Python
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()
|