first commit

ailai_pc/angle.py

@@ -0,0 +1,145 @@
+from ultralytics import YOLO
+import cv2
+import os
+import numpy as np
+
+# 设置类别名称（必须与训练时一致）
+CLASS_NAMES = ['ban', 'bag']  # ✅ 确保顺序正确，对应模型的 class_id
+COLORS = [(0, 255, 0), (255, 0, 0)]  # ban: 绿色, bag: 蓝色
+
+
+def get_best_angles_per_class(image_path, weight_path, return_degree=False):
+    """
+    输入：
+        image_path: 图像路径
+        weight_path: YOLO OBB 权重路径
+        return_degree: 是否返回角度（单位：度），否则为弧度
+    输出：
+        字典：{ class_name: best_angle 或 None }
+    """
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 无法读取图像：{image_path}")
+        return {cls: None for cls in CLASS_NAMES}
+
+    model = YOLO(weight_path)
+    results = model(img, save=False, imgsz=640, conf=0.15, task='obb')
+    result = results[0]
+
+    boxes = result.obb
+    if boxes is None or len(boxes) == 0:
+        print("⚠️ 未检测到任何目标。")
+        return {cls: None for cls in CLASS_NAMES}
+
+    # 提取数据
+    xywhr = boxes.xywhr.cpu().numpy()      # (N, 5) -> cx, cy, w, h, r (弧度)
+    confs = boxes.conf.cpu().numpy()       # (N,)
+    class_ids = boxes.cls.cpu().numpy().astype(int)  # (N,)
+
+    # 初始化结果字典
+    best_angles = {cls: None for cls in CLASS_NAMES}
+
+    # 对每个类别找置信度最高的框
+    for class_id, class_name in enumerate(CLASS_NAMES):
+        mask = (class_ids == class_id)
+        if not np.any(mask):
+            print(f"🟡 未检测到类别: {class_name}")
+            continue
+
+        # 找该类别中置信度最高的
+        idx_in_class = np.argmax(confs[mask])
+        global_idx = np.where(mask)[0][idx_in_class]
+        angle_rad = xywhr[global_idx][4]
+
+        best_angles[class_name] = np.degrees(angle_rad) if return_degree else angle_rad
+
+    return best_angles
+
+
+def save_obb_visual(image_path, weight_path, save_path):
+    """
+    输入：
+        image_path: 图像路径
+        weight_path: YOLO权重路径
+        save_path: 保存带标注图像路径
+    功能：
+        检测所有 OBB，绘制框、类别名、旋转角度，保存图片
+    """
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 无法读取图像：{image_path}")
+        return
+
+    model = YOLO(weight_path)
+    results = model(img, save=False, imgsz=640, conf=0.05, task='obb')
+    result = results[0]
+
+    boxes = result.obb
+    if boxes is None or len(boxes) == 0:
+        print("⚠️ 未检测到任何目标。")
+        # 仍保存原图
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        cv2.imwrite(save_path, img)
+        return
+
+    # 提取信息
+    xywhr = boxes.xywhr.cpu().numpy()
+    confs = boxes.conf.cpu().numpy()
+    class_ids = boxes.cls.cpu().numpy().astype(int)
+
+    # 绘制
+    annotated_img = img.copy()
+    for i in range(len(boxes)):
+        cx, cy, w, h, r = xywhr[i]
+        angle_deg = np.degrees(r)
+        class_id = class_ids[i]
+        class_name = CLASS_NAMES[class_id] if class_id < len(CLASS_NAMES) else f"cls{class_id}"
+        conf = confs[i]
+        color = COLORS[class_id % len(COLORS)] if class_id < len(CLASS_NAMES) else (128, 128, 128)
+
+        # 绘制旋转框
+        rect = ((cx, cy), (w, h), angle_deg)
+        box_pts = cv2.boxPoints(rect).astype(int)
+        cv2.polylines(annotated_img, [box_pts], isClosed=True, color=color, thickness=2)
+
+        # 标注文本：类别 + 置信度 + 角度
+        text = f"{class_name} {conf:.2f} {angle_deg:.1f}°"
+        font_scale = 0.7
+        thickness = 2
+        text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
+
+        # 文本背景
+        cv2.rectangle(annotated_img,
+                      (box_pts[0][0], box_pts[0][1] - text_size[1] - 8),
+                      (box_pts[0][0] + text_size[0], box_pts[0][1] + 2),
+                      color, -1)
+        # 文本
+        cv2.putText(annotated_img, text,
+                    (box_pts[0][0], box_pts[0][1] - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thickness)
+
+    # 保存
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    cv2.imwrite(save_path, annotated_img)
+    print(f"✅ 检测结果已保存至: {save_path}")
+
+
+# ===============================
+# 示例调用
+# ===============================
+if __name__ == "__main__":
+    weight = r"/home/hx/yolo/ultralytics_yolo11-main/runs/train/exp_obb3/weights/best.pt"
+    image = r"/home/hx/yolo/output_masks/2.jpg"
+    save_path = "./inference_results/visualized_2.jpg"
+
+    # 获取每个类别的最佳角度（以度为单位）
+    angles_deg = get_best_angles_per_class(image, weight, return_degree=True)
+    print("\n🎯 各类别最佳旋转角度（度）：")
+    for cls_name, angle in angles_deg.items():
+        if angle is not None:
+            print(f"  {cls_name}: {angle:.2f}°")
+        else:
+            print(f"  {cls_name}: 未检测到")
+
+    # 可视化所有检测结果
+    save_obb_visual(image, weight, save_path)