中交三航控制系统

src/vision/anger_caculate.py

@@ -3,12 +3,13 @@ import os
 import numpy as np
 from ultralytics import YOLO
 
-def predict_obb_best_angle(model=None, model_path=None, image_path=None, save_path=None):
+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_path: 图片路径
+        image: 图像数组（numpy array）
+        image_path: 图片路径（当image为None时使用）
         save_path: 可选，保存带标注图像
     输出：
         angle_deg: 置信度最高两个框的主方向夹角（度），如果检测少于两个目标返回 None
@@ -16,19 +17,22 @@ def predict_obb_best_angle(model=None, model_path=None, image_path=None, save_pa
     """
     # 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. 读取图像
-    img = cv2.imread(image_path)
-    if img is None:
-        print(f"无法读取图像: {image_path}")
-        return None, None
+    # 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')