88 lines
3.1 KiB
Python
88 lines
3.1 KiB
Python
|
import cv2
|
|||
|
|
import os
|
|||
|
|
import numpy as np
|
|||
|
|
from ultralytics import YOLO
|
|||
|
|
|
|||
|
|
def predict_obb_best_angle(model=None, model_path=None, image=None, image_path=None, save_path=None):
|
|||
|
|
"""
|
|||
|
|
输入：
|
|||
|
|
model: 预加载的YOLO模型实例（可选）
|
|||
|
|
model_path: YOLO 权重路径（当model为None时使用）
|
|||
|
|
image: 图像数组（numpy array）
|
|||
|
|
image_path: 图片路径（当image为None时使用）
|
|||
|
|
save_path: 可选，保存带标注图像
|
|||
|
|
输出：
|
|||
|
|
angle_deg: 置信度最高两个框的主方向夹角（度），如果检测少于两个目标返回 None
|
|||
|
|
annotated_img: 可视化图像
|
|||
|
|
"""
|
|||
|
|
# 1. 使用预加载的模型或加载新模型
|
|||
|
|
if model is not None:
|
|||
|
|
loaded_model = model
|
|||
|
|
elif model_path is not None:
|
|||
|
|
loaded_model = YOLO(model_path)
|
|||
|
|
else:
|
|||
|
|
raise ValueError("必须提供model或model_path参数")
|
|||
|
|
|
|||
|
|
# 2. 读取图像（优先使用传入的图像数组）
|
|||
|
|
if image is not None:
|
|||
|
|
img = image
|
|||
|
|
elif image_path is not None:
|
|||
|
|
img = cv2.imread(image_path)
|
|||
|
|
if img is None:
|
|||
|
|
print(f"无法读取图像: {image_path}")
|
|||
|
|
return None, None
|
|||
|
|
else:
|
|||
|
|
raise ValueError("必须提供image或image_path参数")
|
|||
|
|
|
|||
|
|
# 3. 推理 OBB
|
|||
|
|
results = loaded_model(img, save=False, imgsz=640, conf=0.5, mode='obb')
|
|||
|
|
result = results[0]
|
|||
|
|
|
|||
|
|
# 4. 可视化
|
|||
|
|
annotated_img = result.plot()
|
|||
|
|
if save_path:
|
|||
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|||
|
|
cv2.imwrite(save_path, annotated_img)
|
|||
|
|
print(f"推理结果已保存至: {save_path}")
|
|||
|
|
|
|||
|
|
# 5. 提取旋转角度和置信度
|
|||
|
|
boxes = result.obb
|
|||
|
|
if boxes is None or len(boxes) < 2:
|
|||
|
|
print("检测到少于两个目标，无法计算夹角。")
|
|||
|
|
return None, annotated_img
|
|||
|
|
|
|||
|
|
box_info = []
|
|||
|
|
for box in boxes:
|
|||
|
|
conf = box.conf.cpu().numpy()[0]
|
|||
|
|
cx, cy, w, h, r_rad = box.xywhr.cpu().numpy()[0]
|
|||
|
|
direction = r_rad if w >= h else r_rad + np.pi/2
|
|||
|
|
direction = direction % np.pi
|
|||
|
|
box_info.append((conf, direction))
|
|||
|
|
|
|||
|
|
# 6. 取置信度最高两个框
|
|||
|
|
box_info = sorted(box_info, key=lambda x: x[0], reverse=True)
|
|||
|
|
dir1, dir2 = box_info[0][1], box_info[1][1]
|
|||
|
|
|
|||
|
|
# 7. 计算夹角（最小夹角，0~90°）
|
|||
|
|
diff = abs(dir1 - dir2)
|
|||
|
|
diff = min(diff, np.pi - diff)
|
|||
|
|
angle_deg = np.degrees(diff)
|
|||
|
|
|
|||
|
|
print(f"置信度最高两个框主方向夹角: {angle_deg:.2f}°")
|
|||
|
|
return angle_deg, annotated_img
|
|||
|
|
|
|||
|
|
|
|||
|
|
# ------------------- 测试 -------------------
|
|||
|
|
# if __name__ == "__main__":
|
|||
|
|
# weight_path = r'angle.pt'
|
|||
|
|
# image_path = r"./test_image/3.jpg"
|
|||
|
|
# save_path = "./inference_results/detected_3.jpg"
|
|||
|
|
#
|
|||
|
|
# #angle_deg, annotated_img = predict_obb_best_angle(weight_path, image_path, save_path)
|
|||
|
|
# angle_deg,_ = predict_obb_best_angle(model_path=weight_path, image_path=image_path, save_path=save_path)
|
|||
|
|
# annotated_img = None
|
|||
|
|
# print(angle_deg)
|
|||
|
|
# if annotated_img is not None:
|
|||
|
|
# cv2.imshow("YOLO OBB Prediction", annotated_img)
|
|||
|
|
# cv2.waitKey(0)
|
|||
|
|
# cv2.destroyAllWindows()
|