72 lines
2.3 KiB
Python
72 lines
2.3 KiB
Python
# detect_pt.py
|
||
import cv2
|
||
import torch
|
||
from ultralytics import YOLO
|
||
|
||
# ======================
|
||
# 配置参数
|
||
# ======================
|
||
MODEL_PATH = 'best.pt' # 你的训练模型路径（yolov8n.pt 或你自己训练的）
|
||
#IMG_PATH = '/home/hx/开发/ailai_image_obb/ailai_pc/train/192.168.0.234_01_202510141514352.jpg' # 测试图像路径
|
||
IMG_PATH = '1.jpg'
|
||
OUTPUT_PATH = '/home/hx/开发/ailai_image_obb/ailai_pc/output_pt.jpg' # 可视化结果保存路径
|
||
CONF_THRESH = 0.5 # 置信度阈值
|
||
CLASS_NAMES = ['bag'] # 你的类别名列表（按训练时顺序）
|
||
|
||
# 是否显示窗口（适合有 GUI 的 PC）
|
||
SHOW_IMAGE = True
|
||
|
||
# ======================
|
||
# 主函数
|
||
# ======================
|
||
def main():
|
||
# 检查 CUDA
|
||
device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
print(f"✅ 使用设备: {device}")
|
||
|
||
# 加载模型
|
||
print("➡️ 加载 YOLO 模型...")
|
||
model = YOLO(MODEL_PATH) # 自动加载架构和权重
|
||
model.to(device)
|
||
|
||
# 推理
|
||
print("➡️ 开始推理...")
|
||
results = model(IMG_PATH, imgsz=640, conf=CONF_THRESH, device=device)
|
||
|
||
# 获取第一张图的结果
|
||
r = results[0]
|
||
|
||
# 获取原始图像（BGR）
|
||
img = cv2.imread(IMG_PATH)
|
||
if img is None:
|
||
raise FileNotFoundError(f"无法读取图像: {IMG_PATH}")
|
||
|
||
print("\n📋 检测结果:")
|
||
for box in r.boxes:
|
||
# 获取数据
|
||
xyxy = box.xyxy[0].cpu().numpy() # [x1, y1, x2, y2]
|
||
conf = box.conf.cpu().numpy()[0] # 置信度
|
||
cls_id = int(box.cls.cpu().numpy()[0]) # 类别 ID
|
||
cls_name = CLASS_NAMES[cls_id] # 类别名
|
||
|
||
x1, y1, x2, y2 = map(int, xyxy)
|
||
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 SHOW_IMAGE:
|
||
cv2.imshow("YOLOv8 Detection", img)
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows()
|
||
|
||
if __name__ == '__main__':
|
||
main() |