最新推送

angle_base_obb/angle_caculate.py

@@ -1,78 +0,0 @@
-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.3, mode='obb')
-    result = results[0]
-    print(result)
-
-    # 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'obb.pt'
-    #weight_path = r'obb.pt'
-    image_path = r"1.png"
-    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

@@ -6,19 +6,28 @@ 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 推理，计算每张图像检测目标的主方向和夹角。
+def draw_direction(img, cx, cy, angle_deg, length=80, color=(0, 255, 0)):
+    """画主方向箭头"""
+    theta = np.radians(angle_deg)
+    x2 = int(cx + length * np.cos(theta))
+    y2 = int(cy + length * np.sin(theta))
+    cv2.arrowedLine(
+        img,
+        (int(cx), int(cy)),
+        (x2, y2),
+        color,
+        2,
+        tipLength=0.2
+    )
 
-    输入：
-        model_path: YOLO 权重路径
-        image_dir: 图像文件夹路径
-        output_dir: 输出结果保存路径
-        conf_thresh: 置信度阈值
-        imgsz: 输入图像大小
-    输出：
-        results_dict: {image_filename: {'angles_deg': [...], 'pairwise_angles_deg': [...]}}
-    """
+
+def process_obb_images(
+    model_path,
+    image_dir,
+    output_dir="./inference_results",
+    conf_thresh=0.15,
+    imgsz=640
+):
     os.makedirs(output_dir, exist_ok=True)
     results_dict = {}
 
@@ -26,8 +35,11 @@ def process_obb_images(model_path, image_dir, output_dir="./inference_results",
     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]
+    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
@@ -40,63 +52,101 @@ def process_obb_images(model_path, image_dir, output_dir="./inference_results",
 
         img = cv2.imread(img_path)
         if img is None:
-            print(f"❌ 跳过：无法读取图像 {img_path}")
+            print("❌ 读取失败，跳过")
             continue
 
-        # 推理 OBB
-        results = model(img, save=False, imgsz=imgsz, conf=conf_thresh, mode='obb')
+        # ---------- 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 = []
+        centers = []
+
         if boxes is None or len(boxes) == 0:
-            print("❌ 该图像中未检测到任何目标")
+            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}°")
 
-        # 两两夹角
+                angles_deg.append(angle_deg)
+                centers.append((int(cx), int(cy)))
+
+                print(f"  Box {i + 1} 主方向: {angle_deg:.2f} deg")
+
+                # 主方向可视化
+                draw_direction(annotated_img, cx, cy, angle_deg)
+                cv2.circle(annotated_img, (int(cx), int(cy)), 4, (0, 0, 255), -1)
+
+        # ---------- 两两夹角 ----------
         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]))
+                    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}°")
+                    angle_ij = np.degrees(min_diff_rad)
+
+                    pairwise_angles_deg.append(angle_ij)
+                    print(
+                        f"  Box {i + 1} 与 Box {j + 1} 夹角: {angle_ij:.2f} deg"
+                    )
+
+        # ---------- 右上角粗体 angle ----------
+        if pairwise_angles_deg:
+            max_angle = max(pairwise_angles_deg)
+            h, w = annotated_img.shape[:2]
+
+            text = f"angle: {max_angle:.1f} deg"
+
+            # 粗体效果（多次叠加）
+            for dx, dy in [(0, 0), (1, 0), (0, 1), (1, 1)]:
+                cv2.putText(
+                    annotated_img,
+                    text,
+                    (w - 300 + dx, 40 + dy),
+                    cv2.FONT_HERSHEY_SIMPLEX,
+                    1.2,
+                    (0, 0, 255),
+                    3
+                )
+
+        # ---------- 保存 ----------
+        save_path = os.path.join(output_dir, "detected_" + img_filename)
+        cv2.imwrite(save_path, annotated_img)
+        print(f"✅ 保存完成: {save_path}")
 
-        # 保存每张图像结果
         results_dict[img_filename] = {
             "angles_deg": angles_deg,
             "pairwise_angles_deg": pairwise_angles_deg
         }
 
-    print("\n所有图像处理完成！")
+    print("\n🎉 全部处理完成")
     return results_dict
 
 
-# ------------------- 测试调用 -------------------
+# ------------------- 主入口 -------------------
 if __name__ == "__main__":
-    MODEL_PATH = r'/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new3/weights/best.pt'
-    IMAGE_SOURCE_DIR = r"/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb/val"
+    MODEL_PATH = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb_new3/weights/best.pt"
+    IMAGE_SOURCE_DIR = r"/home/hx/yolo/angle_base_obb/test_image"
     OUTPUT_DIR = "./inference_results"
 
-    results = process_obb_images(MODEL_PATH, IMAGE_SOURCE_DIR, OUTPUT_DIR)
+    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']}")
+        print(f"\n{img_name}")
+        print("主方向角:", info["angles_deg"])
+        print("两两夹角:", info["pairwise_angles_deg"])