增加部署代码

angle_base_obb/anger_caculate.py

@@ -4,32 +4,30 @@ import numpy as np
 import os
 
 # 1. 加载模型
-model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb/weights/best.pt')
+model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb2/weights/best.pt')
 
 # 2. 读取图像
-img_path = r"/output_masks/3.png"
+img_path = r"/home/hx/yolo/output_masks/2.jpg"
 img = cv2.imread(img_path)
 
-# ✅ 检查图像是否加载成功
 if img is None:
     print(f"❌ 错误：无法读取图像！请检查路径：{img_path}")
-    print("💡 提示：可能是文件不存在、路径错误或图像损坏")
     exit(1)
 
-# 3. 预测（使用 OBB 模式！）
+# 3. 预测（OBB 模式）
 results = model(
     img,
     save=False,
-    imgsz=1280,      # 必须和训练时一致
-    conf=0.25,
-    mode='obb'       # 启用旋转框模式
+    imgsz=640,
+    conf=0.15,
+    mode='obb'
 )
 
 # 4. 获取结果并绘制
 result = results[0]
-annotated_img = result.plot()  # 自动绘制旋转框
+annotated_img = result.plot()
 
-# ✅ 5. 保存推理结果图像
+# 5. 保存结果
 output_dir = "./inference_results"
 os.makedirs(output_dir, exist_ok=True)
 filename = os.path.basename(img_path)
@@ -37,38 +35,51 @@ save_path = os.path.join(output_dir, "detected_" + filename)
 cv2.imwrite(save_path, annotated_img)
 print(f"✅ 推理结果已保存至: {save_path}")
 
-# ✅ 6. 提取旋转框信息并计算夹角（修正版）
+# 6. 提取旋转框并计算 **两个框之间的夹角**
 boxes = result.obb
 if boxes is None or len(boxes) == 0:
     print("❌ No objects detected.")
 else:
     print(f"✅ Detected {len(boxes)} object(s):")
-    angles = []  # 存储每个框的旋转角度（角度制）
+    directions = []  # 存储每个框的主方向（弧度），归一化到 [0, π)
 
     for i, box in enumerate(boxes):
         cls = int(box.cls.cpu().numpy()[0])
         conf = box.conf.cpu().numpy()[0]
+        xywhr = box.xywhr.cpu().numpy()[0]  # [cx, cy, w, h, r]
+        cx, cy, w, h, r_rad = xywhr
 
-        # ✅ 正确方式：使用 .xywhr 获取旋转角度（新版本 API）
-        xywhr = box.xywhr.cpu().numpy()[0]  # [x_center, y_center, width, height, rotation]
-        angle_rad = xywhr[4]  # 第5个值是旋转角度（弧度）
-        angle_deg = np.degrees(angle_rad)  # 转为角度
+        # 确定主方向（长边方向）
+        if w >= h:
+            direction = r_rad  # 长边方向就是 r
+        else:
+            direction = r_rad + np.pi / 2  # 长边方向是 r + 90°
 
-        angles.append(angle_deg)
-        print(f"  Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, Angle: {angle_deg:.2f}°")
+        # 归一化到 [0, π)
+        direction = direction % np.pi
 
-    # ✅ 计算任意两个框之间的最小夹角差
-    if len(angles) >= 2:
-        print("\n🔍 计算旋转框之间的夹角差：")
-        for i in range(len(angles)):
-            for j in range(i + 1, len(angles)):
-                diff = abs(angles[i] - angles[j])
-                min_angle_diff = min(diff, 180 - diff)  # 取最小夹角（0~180°内）
-                print(f"  Box {i+1} 与 Box {j+1} 的最小夹角差: {min_angle_diff:.2f}°")
+        directions.append(direction)
+        angle_deg = np.degrees(direction)
+        print(f"  Box {i+1}: Class: {cls}, Confidence: {conf:.3f}, 主方向: {angle_deg:.2f}°")
+
+    # ✅ 计算任意两个框之间的夹角（最小夹角，0° ~ 90°）
+    if len(directions) >= 2:
+        print("\n🔍 计算两个旋转框之间的夹角（主方向夹角）：")
+        for i in range(len(directions)):
+            for j in range(i + 1, len(directions)):
+                dir1 = directions[i]
+                dir2 = directions[j]
+
+                # 计算方向差（取最小夹角，考虑周期性）
+                diff = abs(dir1 - dir2)
+                diff = min(diff, np.pi - diff)  # 最小夹角（0 ~ π/2）
+                diff_deg = np.degrees(diff)
+
+                print(f"  Box {i+1} 与 Box {j+1} 之间的夹角: {diff_deg:.2f}°")
     else:
         print("⚠️  检测到少于两个目标，无法计算夹角。")
 
-# ✅ 7. 显示图像
+# 7. 显示图像
 cv2.imshow("YOLO OBB Prediction", annotated_img)
 cv2.waitKey(0)
 cv2.destroyAllWindows()