106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
|
from ultralytics import YOLO
|
|||
|
|
import cv2
|
|||
|
|
import os
|
|||
|
|
import numpy as np
|
|||
|
|
|
|||
|
|
def get_best_obb_angle(image_path, weight_path, return_degree=False):
|
|||
|
|
"""
|
|||
|
|
输入：
|
|||
|
|
image_path: 图像路径
|
|||
|
|
weight_path: YOLO权重路径
|
|||
|
|
return_degree: 是否返回角度单位为度，默认 False（返回弧度）
|
|||
|
|
输出：
|
|||
|
|
置信度最高目标的旋转角
|
|||
|
|
如果未检测到目标返回 None
|
|||
|
|
"""
|
|||
|
|
# 读取图像
|
|||
|
|
img = cv2.imread(image_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"❌ 无法读取图像：{image_path}")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 加载模型并预测
|
|||
|
|
model = YOLO(weight_path)
|
|||
|
|
results = model(img, save=False, imgsz=640, conf=0.15, mode='obb')
|
|||
|
|
result = results[0]
|
|||
|
|
|
|||
|
|
boxes = result.obb
|
|||
|
|
if not boxes:
|
|||
|
|
print("⚠️ 未检测到目标。")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 取置信度最高框的旋转角
|
|||
|
|
best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0])
|
|||
|
|
r = best_box.xywhr.cpu().numpy()[0][4] # 弧度
|
|||
|
|
|
|||
|
|
if return_degree:
|
|||
|
|
return np.degrees(r)
|
|||
|
|
else:
|
|||
|
|
return r
|
|||
|
|
|
|||
|
|
|
|||
|
|
def save_obb_visual(image_path, weight_path, save_path):
|
|||
|
|
"""
|
|||
|
|
输入：
|
|||
|
|
image_path: 图像路径
|
|||
|
|
weight_path: YOLO权重路径
|
|||
|
|
save_path: 保存带角度标注图像路径
|
|||
|
|
功能：
|
|||
|
|
检测 OBB 并标注置信度最高框旋转角度，保存图片
|
|||
|
|
"""
|
|||
|
|
img = cv2.imread(image_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"❌ 无法读取图像：{image_path}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
model = YOLO(weight_path)
|
|||
|
|
results = model(img, save=False, imgsz=640, conf=0.15, mode='obb')
|
|||
|
|
result = results[0]
|
|||
|
|
|
|||
|
|
boxes = result.obb
|
|||
|
|
if not boxes:
|
|||
|
|
print("⚠️ 未检测到目标。")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0])
|
|||
|
|
cx, cy, w, h, r = best_box.xywhr.cpu().numpy()[0]
|
|||
|
|
angle_deg = np.degrees(r)
|
|||
|
|
|
|||
|
|
# 绘制 OBB
|
|||
|
|
annotated_img = img.copy()
|
|||
|
|
rect = ((cx, cy), (w, h), angle_deg)
|
|||
|
|
box_pts = cv2.boxPoints(rect).astype(int)
|
|||
|
|
cv2.polylines(annotated_img, [box_pts], isClosed=True, color=(0, 255, 0), thickness=2)
|
|||
|
|
|
|||
|
|
# 标注角度
|
|||
|
|
text = f"{angle_deg:.1f}°"
|
|||
|
|
font_scale = max(0.5, min(w, h)/100)
|
|||
|
|
thickness = 2
|
|||
|
|
text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
|
|||
|
|
text_x = int(cx - text_size[0]/2)
|
|||
|
|
text_y = int(cy + text_size[1]/2)
|
|||
|
|
cv2.putText(annotated_img, text, (text_x, text_y),
|
|||
|
|
cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 255), thickness)
|
|||
|
|
|
|||
|
|
# 保存
|
|||
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|||
|
|
cv2.imwrite(save_path, annotated_img)
|
|||
|
|
print(f"✅ 检测结果已保存至: {save_path}")
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ===============================
|
|||
|
|
# 示例调用
|
|||
|
|
# ===============================
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
weight = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb3/weights/best.pt"
|
|||
|
|
image = r"/home/hx/yolo/output_masks/2.jpg"
|
|||
|
|
save_path = "./inference_results/best_detected_2.jpg"
|
|||
|
|
|
|||
|
|
angle_rad = get_best_obb_angle(image, weight)
|
|||
|
|
print(f"旋转角（弧度）：{angle_rad:.4f}")
|
|||
|
|
|
|||
|
|
angle_deg = get_best_obb_angle(image, weight, return_degree=True)
|
|||
|
|
print(f"旋转角（度）：{angle_deg:.2f}°")
|
|||
|
|
|
|||
|
|
save_obb_visual(image, weight, save_path)
|