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

angle_base_obb/anger_caculate.py

@@ -1,85 +0,0 @@
-from ultralytics import YOLO
-import cv2
-import numpy as np
-import os
-
-# 1. 加载模型
-model = YOLO(r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb2/weights/best.pt')
-
-# 2. 读取图像
-img_path = r"/home/hx/yolo/output_masks/2.jpg"
-img = cv2.imread(img_path)
-
-if img is None:
-    print(f"❌ 错误：无法读取图像！请检查路径：{img_path}")
-    exit(1)
-
-# 3. 预测（OBB 模式）
-results = model(
-    img,
-    save=False,
-    imgsz=640,
-    conf=0.15,
-    mode='obb'
-)
-
-# 4. 获取结果并绘制
-result = results[0]
-annotated_img = result.plot()
-
-# 5. 保存结果
-output_dir = "./inference_results"
-os.makedirs(output_dir, exist_ok=True)
-filename = os.path.basename(img_path)
-save_path = os.path.join(output_dir, "detected_" + filename)
-cv2.imwrite(save_path, annotated_img)
-print(f"✅ 推理结果已保存至: {save_path}")
-
-# 6. 提取旋转框并计算 **两个框之间的夹角**
-boxes = result.obb
-if boxes is None or len(boxes) == 0:
-    print("❌ No objects detected.")
-else:
-    print(f"✅ Detected {len(boxes)} object(s):")
-    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
-
-        # 确定主方向（长边方向）
-        if w >= h:
-            direction = r_rad  # 长边方向就是 r
-        else:
-            direction = r_rad + np.pi / 2  # 长边方向是 r + 90°
-
-        # 归一化到 [0, π)
-        direction = direction % np.pi
-
-        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. 显示图像
-cv2.imshow("YOLO OBB Prediction", annotated_img)
-cv2.waitKey(0)
-cv2.destroyAllWindows()

angle_base_obb/anger_caculate_file.py

@@ -1,103 +0,0 @@
-from ultralytics import YOLO
-import cv2
-import numpy as np
-import os
-
-# ================== 配置参数 ==================
-MODEL_PATH = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb4/weights/best.pt"
-IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb2/test"  # 图像文件夹路径
-OUTPUT_DIR = "./inference_results"                # 输出结果保存路径
-IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}
-
-# 创建输出目录
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-# 1. 加载模型
-print("🔄 加载 YOLO 模型...")
-model = YOLO(MODEL_PATH)
-print("✅ 模型加载完成")
-
-# 获取所有图像文件
-image_files = [
-    f for f in os.listdir(IMAGE_SOURCE_DIR)
-    if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS
-]
-
-if not image_files:
-    print(f"❌ 错误：在路径中未找到图像文件：{IMAGE_SOURCE_DIR}")
-    exit(1)
-
-print(f"📁 发现 {len(image_files)} 张图像待处理")
-
-# ================== 批量处理每张图像 ==================
-for img_filename in image_files:
-    img_path = os.path.join(IMAGE_SOURCE_DIR, img_filename)
-    print(f"\n🖼️  正在处理：{img_filename}")
-
-    # 读取图像
-    img = cv2.imread(img_path)
-    if img is None:
-        print(f"❌ 跳过：无法读取图像 {img_path}")
-        continue
-
-    # 推理（OBB 模式）
-    results = model(
-        img,
-        save=False,
-        imgsz=640,
-        conf=0.15,
-        mode='obb'
-    )
-
-    result = results[0]
-    annotated_img = result.plot()  # 绘制旋转框
-
-    # 保存结果图像
-    save_path = os.path.join(OUTPUT_DIR, "detected_" + img_filename)
-    cv2.imwrite(save_path, annotated_img)
-    print(f"✅ 推理结果已保存至: {save_path}")
-
-    # 提取旋转框信息
-    boxes = result.obb
-    directions = []  # 存储每个框的主方向（弧度），归一化到 [0, π)
-
-    if boxes is None or len(boxes) == 0:
-        print("❌ 该图像中未检测到任何目标")
-    else:
-        print(f"✅ 检测到 {len(boxes)} 个目标：")
-        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
-
-            # 确定主方向（长边方向）
-            if w >= h:
-                direction = r_rad  # 长边方向
-            else:
-                direction = r_rad + np.pi / 2  # 长边是宽的方向
-
-            # 归一化到 [0, π)
-            direction = direction % np.pi
-            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)
-                    min_diff_rad = min(diff, np.pi - diff)  # 最小夹角（考虑周期性）
-                    min_diff_deg = np.degrees(min_diff_rad)
-
-                    print(f"  Box {i+1} 与 Box {j+1} 之间夹角: {min_diff_deg:.2f}°")
-        else:
-            print("⚠️  检测到少于两个目标，无法计算夹角。")
-
-print("\n🎉 所有图像处理完成！")

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()

angle_base_obb/angle_caculate_file.py

@@ -0,0 +1,102 @@
+import cv2
+import os
+import numpy as np
+from ultralytics import YOLO
+
+IMG_EXTENSIONS = {'.jpg', '.jpeg', '.png', '.bmp', '.tif', '.tiff', '.webp'}
+
+
+def process_obb_images(model_path, image_dir, output_dir="./inference_results", conf_thresh=0.15, imgsz=640):
+    """
+    批量处理图像的 OBB 推理，计算每张图像检测目标的主方向和夹角。
+
+    输入：
+        model_path: YOLO 权重路径
+        image_dir: 图像文件夹路径
+        output_dir: 输出结果保存路径
+        conf_thresh: 置信度阈值
+        imgsz: 输入图像大小
+    输出：
+        results_dict: {image_filename: {'angles_deg': [...], 'pairwise_angles_deg': [...]}}
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    results_dict = {}
+
+    print("加载 YOLO 模型...")
+    model = YOLO(model_path)
+    print("✅ 模型加载完成")
+
+    # 获取图像文件
+    image_files = [f for f in os.listdir(image_dir) if os.path.splitext(f.lower())[1] in IMG_EXTENSIONS]
+    if not image_files:
+        print(f"❌ 未找到图像文件：{image_dir}")
+        return results_dict
+
+    print(f"发现 {len(image_files)} 张图像待处理")
+
+    for img_filename in image_files:
+        img_path = os.path.join(image_dir, img_filename)
+        print(f"\n正在处理：{img_filename}")
+
+        img = cv2.imread(img_path)
+        if img is None:
+            print(f"❌ 跳过：无法读取图像 {img_path}")
+            continue
+
+        # 推理 OBB
+        results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode='obb')
+        result = results[0]
+        annotated_img = result.plot()
+
+        # 保存可视化
+        save_path = os.path.join(output_dir, "detected_" + img_filename)
+        cv2.imwrite(save_path, annotated_img)
+        print(f"✅ 推理结果已保存至: {save_path}")
+
+        # 提取旋转角
+        boxes = result.obb
+        angles_deg = []
+        if boxes is None or len(boxes) == 0:
+            print("❌ 该图像中未检测到任何目标")
+        else:
+            for i, box in enumerate(boxes):
+                cls = int(box.cls.cpu().numpy()[0])
+                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
+                angle_deg = np.degrees(direction)
+                angles_deg.append(angle_deg)
+                print(f"  Box {i + 1}: Class={cls}, Conf={conf:.3f}, 主方向={angle_deg:.2f}°")
+
+        # 两两夹角
+        pairwise_angles_deg = []
+        if len(angles_deg) >= 2:
+            for i in range(len(angles_deg)):
+                for j in range(i + 1, len(angles_deg)):
+                    diff_rad = abs(np.radians(angles_deg[i]) - np.radians(angles_deg[j]))
+                    min_diff_rad = min(diff_rad, np.pi - diff_rad)
+                    pairwise_angles_deg.append(np.degrees(min_diff_rad))
+                    print(f"  Box {i + 1} 与 Box {j + 1} 夹角: {np.degrees(min_diff_rad):.2f}°")
+
+        # 保存每张图像结果
+        results_dict[img_filename] = {
+            "angles_deg": angles_deg,
+            "pairwise_angles_deg": pairwise_angles_deg
+        }
+
+    print("\n所有图像处理完成！")
+    return results_dict
+
+
+# ------------------- 测试调用 -------------------
+if __name__ == "__main__":
+    MODEL_PATH = r'best.pt'
+    IMAGE_SOURCE_DIR = r"./test_image"
+    OUTPUT_DIR = "./inference_results"
+
+    results = process_obb_images(MODEL_PATH, IMAGE_SOURCE_DIR, OUTPUT_DIR)
+    for img_name, info in results.items():
+        print(f"\n {img_name}:")
+        print(f"主方向角度列表: {info['angles_deg']}")
+        print(f"两两夹角列表: {info['pairwise_angles_deg']}")