76 lines
2.3 KiB
Python
76 lines
2.3 KiB
Python
from ultralytics import YOLO
|
||
from ultralytics.utils.ops import non_max_suppression
|
||
import torch
|
||
import cv2
|
||
|
||
# ======================
|
||
# 配置参数
|
||
# ======================
|
||
MODEL_PATH = '/home/hx/开发/ailai_image_obb/ailai_pc/best12.pt'
|
||
IMG_PATH = '1.jpg'
|
||
OUTPUT_PATH = 'output_pt.jpg'
|
||
CONF_THRESH = 0.5
|
||
IOU_THRESH = 0.45
|
||
CLASS_NAMES = ['bag']
|
||
|
||
# ======================
|
||
# 主函数（优化版）
|
||
# ======================
|
||
def main():
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
print(f"✅ 使用设备: {device}")
|
||
|
||
# 加载模型
|
||
model = YOLO(MODEL_PATH)
|
||
model.to(device)
|
||
|
||
# 推理：获取原始结果（不立即解析）
|
||
print("➡️ 开始推理...")
|
||
results = model(IMG_PATH, imgsz=640, conf=CONF_THRESH, device=device, verbose=True)
|
||
|
||
# 获取第一张图的结果
|
||
r = results[0]
|
||
|
||
# 🚀 关键：使用原始 tensor 在 GPU 上处理
|
||
# pred: [x1, y1, x2, y2, conf, cls] 形状为 [num_boxes, 6]
|
||
pred = r.boxes.data # 已经在 GPU 上，类型: torch.Tensor
|
||
|
||
# 🔍 在 GPU 上做 NMS（这才是正确姿势）
|
||
# 注意：non_max_suppression 输入是 [batch, num_boxes, 6]
|
||
det = non_max_suppression(
|
||
pred.unsqueeze(0), # 增加 batch 维度
|
||
conf_thres=CONF_THRESH,
|
||
iou_thres=IOU_THRESH,
|
||
classes=None,
|
||
agnostic=False,
|
||
max_det=100
|
||
)[0] # 取第一个（也是唯一一个）batch
|
||
|
||
# ✅ 此时所有后处理已完成，现在才从 GPU 拷贝到 CPU
|
||
if det is not None and len(det):
|
||
det = det.cpu().numpy() # ← 只拷贝一次！
|
||
else:
|
||
det = []
|
||
|
||
# 读取图像
|
||
img = cv2.imread(IMG_PATH)
|
||
if img is None:
|
||
raise FileNotFoundError(f"无法读取图像: {IMG_PATH}")
|
||
|
||
print("\n📋 检测结果:")
|
||
for *xyxy, conf, cls_id in det:
|
||
x1, y1, x2, y2 = map(int, xyxy)
|
||
cls_name = CLASS_NAMES[int(cls_id)]
|
||
print(f" 类别: {cls_name}, 置信度: {conf:.3f}, 框: [{x1}, {y1}, {x2}, {y2}]")
|
||
|
||
# 画框和标签
|
||
cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
||
label = f"{cls_name} {conf:.2f}"
|
||
cv2.putText(img, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
|
||
|
||
# 保存结果
|
||
cv2.imwrite(OUTPUT_PATH, img)
|
||
print(f"\n🖼️ 可视化结果已保存: {OUTPUT_PATH}")
|
||
|
||
if __name__ == '__main__':
|
||
main() |