增加ailai的旋转检测的推理和部署

angle_base_obb/angle_caculate.py

@@ -0,0 +1,76 @@
+import cv2
+import os
+import numpy as np
+from ultralytics import YOLO
+
+def predict_obb_best_angle(model_path, image_path, save_path=None):
+    """
+    输入：
+        model_path: YOLO 权重路径
+        image_path: 图片路径
+        save_path: 可选，保存带标注图像
+    输出：
+        angle_deg: 置信度最高两个框的主方向夹角（度），如果检测少于两个目标返回 None
+        annotated_img: 可视化图像
+    """
+    # 1. 加载模型
+    model = YOLO(model_path)
+
+    # 2. 读取图像
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"无法读取图像: {image_path}")
+        return None, None
+
+    # 3. 推理 OBB
+    results = 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'best.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(weight_path, image_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()