Merge remote-tracking branch 'origin/main'

ailai_obb/angle.py

@@ -0,0 +1,105 @@
+from ultralytics import YOLO
+import cv2
+import os
+import numpy as np
+
+def get_best_obb_angle(image_path, weight_path, return_degree=False):
+    """
+    输入：
+        image_path: 图像路径
+        weight_path: YOLO权重路径
+        return_degree: 是否返回角度单位为度，默认 False（返回弧度）
+    输出：
+        置信度最高目标的旋转角
+        如果未检测到目标返回 None
+    """
+    # 读取图像
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 无法读取图像：{image_path}")
+        return None
+
+    # 加载模型并预测
+    model = YOLO(weight_path)
+    results = model(img, save=False, imgsz=640, conf=0.15, mode='obb')
+    result = results[0]
+
+    boxes = result.obb
+    if not boxes:
+        print("⚠️ 未检测到目标。")
+        return None
+
+    # 取置信度最高框的旋转角
+    best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0])
+    r = best_box.xywhr.cpu().numpy()[0][4]  # 弧度
+
+    if return_degree:
+        return np.degrees(r)
+    else:
+        return r
+
+
+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.15, mode='obb')
+    result = results[0]
+
+    boxes = result.obb
+    if not boxes:
+        print("⚠️ 未检测到目标。")
+        return
+
+    best_box = max(boxes, key=lambda x: x.conf.cpu().numpy()[0])
+    cx, cy, w, h, r = best_box.xywhr.cpu().numpy()[0]
+    angle_deg = np.degrees(r)
+
+    # 绘制 OBB
+    annotated_img = img.copy()
+    rect = ((cx, cy), (w, h), angle_deg)
+    box_pts = cv2.boxPoints(rect).astype(int)
+    cv2.polylines(annotated_img, [box_pts], isClosed=True, color=(0, 255, 0), thickness=2)
+
+    # 标注角度
+    text = f"{angle_deg:.1f}°"
+    font_scale = max(0.5, min(w, h)/100)
+    thickness = 2
+    text_size, _ = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, font_scale, thickness)
+    text_x = int(cx - text_size[0]/2)
+    text_y = int(cy + text_size[1]/2)
+    cv2.putText(annotated_img, text, (text_x, text_y),
+                cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 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/best_detected_2.jpg"
+
+    angle_rad = get_best_obb_angle(image, weight)
+    print(f"旋转角（弧度）：{angle_rad:.4f}")
+
+    angle_deg = get_best_obb_angle(image, weight, return_degree=True)
+    print(f"旋转角（度）：{angle_deg:.2f}°")
+
+    save_obb_visual(image, weight, save_path)

ailai_obb/bag_bushu.py

@@ -0,0 +1,171 @@
+import cv2
+import numpy as np
+import math
+from shapely.geometry import Polygon
+from rknnlite.api import RKNNLite
+import os
+
+CLASSES = ['clamp']
+nmsThresh = 0.4
+objectThresh = 0.5
+
+# ------------------- 工具函数 -------------------
+def letterbox_resize(image, size, bg_color=114):
+    target_width, target_height = size
+    image_height, image_width, _ = image.shape
+    scale = min(target_width / image_width, target_height / image_height)
+    new_width, new_height = int(image_width * scale), int(image_height * scale)
+    image = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
+    offset_x, offset_y = (target_width - new_width) // 2, (target_height - new_height) // 2
+    canvas[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image
+    return canvas, scale, offset_x, offset_y
+
+class DetectBox:
+    def __init__(self, classId, score, xmin, ymin, xmax, ymax, angle):
+        self.classId = classId
+        self.score = score
+        self.xmin = xmin
+        self.ymin = ymin
+        self.xmax = xmax
+        self.ymax = ymax
+        self.angle = angle
+
+def rotate_rectangle(x1, y1, x2, y2, a):
+    cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
+    x1_new = int((x1 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx)
+    y1_new = int((x1 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy)
+    x2_new = int((x2 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx)
+    y2_new = int((x2 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy)
+    x3_new = int((x1 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx)
+    y3_new = int((x1 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy)
+    x4_new = int((x2 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx)
+    y4_new = int((x2 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy)
+    return [(x1_new, y1_new), (x3_new, y3_new), (x2_new, y2_new), (x4_new, y4_new)]
+
+def intersection(g, p):
+    g = Polygon(np.array(g).reshape(-1,2))
+    p = Polygon(np.array(p).reshape(-1,2))
+    if not g.is_valid or not p.is_valid:
+        return 0
+    inter = g.intersection(p).area
+    union = g.area + p.area - inter
+    return 0 if union == 0 else inter / union
+
+def NMS(detectResult):
+    predBoxs = []
+    sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True)
+    for i in range(len(sort_detectboxs)):
+        if sort_detectboxs[i].classId == -1:
+            continue
+        p1 = rotate_rectangle(sort_detectboxs[i].xmin, sort_detectboxs[i].ymin,
+                              sort_detectboxs[i].xmax, sort_detectboxs[i].ymax,
+                              sort_detectboxs[i].angle)
+        predBoxs.append(sort_detectboxs[i])
+        for j in range(i + 1, len(sort_detectboxs)):
+            if sort_detectboxs[j].classId == sort_detectboxs[i].classId:
+                p2 = rotate_rectangle(sort_detectboxs[j].xmin, sort_detectboxs[j].ymin,
+                                      sort_detectboxs[j].xmax, sort_detectboxs[j].ymax,
+                                      sort_detectboxs[j].angle)
+                if intersection(p1, p2) > nmsThresh:
+                    sort_detectboxs[j].classId = -1
+    return predBoxs
+
+def sigmoid(x):
+    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))
+
+def softmax(x, axis=-1):
+    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
+    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
+
+def process(out, model_w, model_h, stride, angle_feature, index, scale_w=1, scale_h=1):
+    class_num = len(CLASSES)
+    angle_feature = angle_feature.reshape(-1)
+    xywh = out[:, :64, :]
+    conf = sigmoid(out[:, 64:, :]).reshape(-1)
+    boxes = []
+    for ik in range(model_h * model_w * class_num):
+        if conf[ik] > objectThresh:
+            w = ik % model_w
+            h = (ik % (model_w * model_h)) // model_w
+            c = ik // (model_w * model_h)
+            # 解析xywh
+            xywh_ = xywh[0, :, (h * model_w) + w].reshape(1, 4, 16, 1)
+            data = np.arange(16).reshape(1, 1, 16, 1)
+            xywh_ = softmax(xywh_, 2)
+            xywh_ = np.sum(xywh_ * data, axis=2).reshape(-1)
+            xywh_add = xywh_[:2] + xywh_[2:]
+            xywh_sub = (xywh_[2:] - xywh_[:2]) / 2
+            # 安全取角度
+            angle_idx = min(index + (h * model_w) + w, len(angle_feature) - 1)
+            angle = (angle_feature[angle_idx] - 0.25) * math.pi
+            cos_a, sin_a = math.cos(angle), math.sin(angle)
+            xy = xywh_sub[0] * cos_a - xywh_sub[1] * sin_a, xywh_sub[0] * sin_a + xywh_sub[1] * cos_a
+            xywh1 = np.array([xy[0] + w + 0.5, xy[1] + h + 0.5, xywh_add[0], xywh_add[1]])
+            xywh1 *= stride
+            xmin = (xywh1[0] - xywh1[2]/2) * scale_w
+            ymin = (xywh1[1] - xywh1[3]/2) * scale_h
+            xmax = (xywh1[0] + xywh1[2]/2) * scale_w
+            ymax = (xywh1[1] + xywh1[3]/2) * scale_h
+            boxes.append(DetectBox(c, conf[ik], xmin, ymin, xmax, ymax, angle))
+    return boxes
+
+# ------------------- 主函数 -------------------
+def detect_boxes_angle_rknn(model_path, image_path, save_path=None):
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 无法读取图像: {image_path}")
+        return None, None
+
+    img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640))
+    infer_img = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
+    infer_img = np.expand_dims(infer_img, 0)
+
+    rknn_lite = RKNNLite(verbose=False)
+    rknn_lite.load_rknn(model_path)
+    rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
+
+    results = rknn_lite.inference([infer_img])
+    detect_boxes = []
+    for x in results[:-1]:
+        index, stride = 0, 0
+        if x.shape[2] == 20:
+            stride, index = 32, 20*4*20*4 + 20*2*20*2
+        elif x.shape[2] == 40:
+            stride, index = 16, 20*4*20*4
+        elif x.shape[2] == 80:
+            stride, index = 8, 0
+        feature = x.reshape(1, 65, -1)
+        detect_boxes += process(feature, x.shape[3], x.shape[2], stride, results[-1], index)
+
+    detect_boxes = NMS(detect_boxes)
+
+    # 输出每个检测框角度
+    for i, box in enumerate(detect_boxes):
+        print(f"框 {i+1}: angle = {box.angle:.4f} rad ({np.degrees(box.angle):.2f}°)")
+        if save_path:
+            xmin = int((box.xmin - offset_x)/scale)
+            ymin = int((box.ymin - offset_y)/scale)
+            xmax = int((box.xmax - offset_x)/scale)
+            ymax = int((box.ymax - offset_y)/scale)
+            points = rotate_rectangle(xmin, ymin, xmax, ymax, box.angle)
+            cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 1)
+            cv2.putText(img, f"{np.degrees(box.angle):.1f}°", (xmin, ymin-5),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1)
+
+    if save_path:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        cv2.imwrite(save_path, img)
+        print(f"✅ 带角度的检测结果已保存到 {save_path}")
+
+    rknn_lite.release()
+    return detect_boxes, img
+
+# ------------------- 使用示例 -------------------
+if __name__ == "__main__":
+    model_path = "obb.rknn"
+    image_path = "2.jpg"
+    save_path = "./inference_results/boxes_with_angle.jpg"
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    detect_boxes_angle_rknn(model_path, image_path, save_path)
+

ailai_obb/bushu_angle.py

@@ -0,0 +1,197 @@
+
+import cv2
+import numpy as np
+import math
+from shapely.geometry import Polygon
+from rknnlite.api import RKNNLite
+import os
+
+# ------------------- 配置 -------------------
+CLASSES = ['clamp']
+nmsThresh = 0.4
+objectThresh = 0.5
+
+# ------------------- 全局原图尺寸 -------------------
+ORIG_W = 2560   # 原图宽
+ORIG_H = 1440   # 原图高
+
+# ------------------- 工具函数 -------------------
+def letterbox_resize(image, size, bg_color=114):
+    target_width, target_height = size
+    image_height, image_width, _ = image.shape
+    scale = min(target_width / image_width, target_height / image_height)
+    new_width, new_height = int(image_width * scale), int(image_height * scale)
+    image_resized = cv2.resize(image, (new_width, new_height), interpolation=cv2.INTER_AREA)
+    canvas = np.ones((target_height, target_width, 3), dtype=np.uint8) * bg_color
+    offset_x, offset_y = (target_width - new_width) // 2, (target_height - new_height) // 2
+    canvas[offset_y:offset_y + new_height, offset_x:offset_x + new_width] = image_resized
+    return canvas, scale, offset_x, offset_y
+
+class DetectBox:
+    def __init__(self, classId, score, xmin, ymin, xmax, ymax, angle):
+        self.classId = classId
+        self.score = score
+        self.xmin = xmin
+        self.ymin = ymin
+        self.xmax = xmax
+        self.ymax = ymax
+        self.angle = angle
+
+def rotate_rectangle(x1, y1, x2, y2, a):
+    cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
+    x1_new = int((x1 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx)
+    y1_new = int((x1 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy)
+    x2_new = int((x2 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx)
+    y2_new = int((x2 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy)
+    x3_new = int((x1 - cx) * math.cos(a) - (y2 - cy) * math.sin(a) + cx)
+    y3_new = int((x1 - cx) * math.sin(a) + (y2 - cy) * math.cos(a) + cy)
+    x4_new = int((x2 - cx) * math.cos(a) - (y1 - cy) * math.sin(a) + cx)
+    y4_new = int((x2 - cx) * math.sin(a) + (y1 - cy) * math.cos(a) + cy)
+    return [(x1_new, y1_new), (x3_new, y3_new), (x2_new, y2_new), (x4_new, y4_new)]
+
+def intersection(g, p):
+    g = Polygon(np.array(g).reshape(-1,2))
+    p = Polygon(np.array(p).reshape(-1,2))
+    if not g.is_valid or not p.is_valid:
+        return 0
+    inter = g.intersection(p).area
+    union = g.area + p.area - inter
+    return 0 if union == 0 else inter / union
+
+def NMS(detectResult):
+    predBoxs = []
+    sort_detectboxs = sorted(detectResult, key=lambda x: x.score, reverse=True)
+    for i in range(len(sort_detectboxs)):
+        if sort_detectboxs[i].classId == -1:
+            continue
+        p1 = rotate_rectangle(sort_detectboxs[i].xmin, sort_detectboxs[i].ymin,
+                              sort_detectboxs[i].xmax, sort_detectboxs[i].ymax,
+                              sort_detectboxs[i].angle)
+        predBoxs.append(sort_detectboxs[i])
+        for j in range(i + 1, len(sort_detectboxs)):
+            if sort_detectboxs[j].classId == sort_detectboxs[i].classId:
+                p2 = rotate_rectangle(sort_detectboxs[j].xmin, sort_detectboxs[j].ymin,
+                                      sort_detectboxs[j].xmax, sort_detectboxs[j].ymax,
+                                      sort_detectboxs[j].angle)
+                if intersection(p1, p2) > nmsThresh:
+                    sort_detectboxs[j].classId = -1
+    return predBoxs
+
+def sigmoid(x):
+    return np.where(x >= 0, 1 / (1 + np.exp(-x)), np.exp(x) / (1 + np.exp(x)))
+
+def softmax(x, axis=-1):
+    exp_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
+    return exp_x / np.sum(exp_x, axis=axis, keepdims=True)
+
+# ------------------- 关键修改：process函数加入scale -------------------
+def process(out, model_w, model_h, stride, angle_feature, index, scale=1.0, offset_x=0, offset_y=0):
+    class_num = len(CLASSES)
+    angle_feature = angle_feature.reshape(-1)
+    xywh = out[:, :64, :]
+    conf = sigmoid(out[:, 64:, :]).reshape(-1)
+    boxes = []
+    for ik in range(model_h * model_w * class_num):
+        if conf[ik] > objectThresh:
+            w = ik % model_w
+            h = (ik % (model_w * model_h)) // model_w
+            c = ik // (model_w * model_h)
+            # 解析xywh
+            xywh_ = xywh[0, :, (h * model_w) + w].reshape(1, 4, 16, 1)
+            xywh_ = softmax(xywh_, 2)
+            data = np.arange(16).reshape(1, 1, 16, 1)
+            xywh_ = np.sum(xywh_ * data, axis=2).reshape(-1)
+            xywh_add = xywh_[:2] + xywh_[2:]
+            xywh_sub = (xywh_[2:] - xywh_[:2]) / 2
+            # 取角度
+            angle_idx = min(index + (h * model_w) + w, len(angle_feature) - 1)
+            angle = (angle_feature[angle_idx] - 0.25) * math.pi
+            cos_a, sin_a = math.cos(angle), math.sin(angle)
+            xy = xywh_sub[0] * cos_a - xywh_sub[1] * sin_a, xywh_sub[0] * sin_a + xywh_sub[1] * cos_a
+            xywh1 = np.array([xy[0] + w + 0.5, xy[1] + h + 0.5, xywh_add[0], xywh_add[1]])
+            xywh1 *= stride
+            # 映射回原图坐标
+            xmin = (xywh1[0] - xywh1[2]/2 - offset_x) / scale
+            ymin = (xywh1[1] - xywh1[3]/2 - offset_y) / scale
+            xmax = (xywh1[0] + xywh1[2]/2 - offset_x) / scale
+            ymax = (xywh1[1] + xywh1[3]/2 - offset_y) / scale
+            boxes.append(DetectBox(c, conf[ik], xmin, ymin, xmax, ymax, angle))
+    return boxes
+
+# ------------------- 新可调用函数 -------------------
+def detect_boxes_rknn(model_path, image_path):
+    img = cv2.imread(image_path)
+    if img is None:
+        print(f"❌ 无法读取图像: {image_path}")
+        return None, None
+
+    img_resized, scale, offset_x, offset_y = letterbox_resize(img, (640, 640))
+    infer_img = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
+    infer_img = np.expand_dims(infer_img, 0)
+
+    rknn_lite = RKNNLite(verbose=False)
+    rknn_lite.load_rknn(model_path)
+    rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0)
+
+    results = rknn_lite.inference([infer_img])
+    detect_boxes = []
+    for x in results[:-1]:
+        index, stride = 0, 0
+        if x.shape[2] == 20:
+            stride, index = 32, 20*4*20*4 + 20*2*20*2
+        elif x.shape[2] == 40:
+            stride, index = 16, 20*4*20*4
+        elif x.shape[2] == 80:
+            stride, index = 8, 0
+        feature = x.reshape(1, 65, -1)
+        detect_boxes += process(feature, x.shape[3], x.shape[2], stride, results[-1], index,
+                                scale=scale, offset_x=offset_x, offset_y=offset_y)
+
+    detect_boxes = NMS(detect_boxes)
+    rknn_lite.release()
+    return detect_boxes, img
+
+# ------------------- 绘制与辅助函数 -------------------
+def get_angles(detect_boxes):
+    return [box.angle for box in detect_boxes]
+
+def draw_boxes(img, detect_boxes, save_path=None):
+    for box in detect_boxes:
+        points = rotate_rectangle(box.xmin, box.ymin, box.xmax, box.ymax, box.angle)
+        cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 1)
+        cv2.putText(img, f"{np.degrees(box.angle):.1f}°", (int(box.xmin), int(box.ymin)-5),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,255), 1)
+    if save_path:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        cv2.imwrite(save_path, img)
+        print(f"✅ 带角度的检测结果已保存到 {save_path}")
+    return img
+
+def visualize_top_box(img, detect_boxes, save_path=None):
+    if not detect_boxes:
+        return img
+    top_box = max(detect_boxes, key=lambda x: x.score)
+    points = rotate_rectangle(top_box.xmin, top_box.ymin, top_box.xmax, top_box.ymax, top_box.angle)
+    cv2.polylines(img, [np.array(points, np.int32)], True, (0, 255, 0), 2)
+    cv2.putText(img, f"{np.degrees(top_box.angle):.1f}°", (int(top_box.xmin), int(top_box.ymin)-5),
+                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,255), 2)
+    if save_path:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        cv2.imwrite(save_path, img)
+    return img
+
+# ------------------- 使用示例 -------------------
+if __name__ == "__main__":
+    model_path = "obb.rknn"
+    image_path = "2.jpg"
+
+    detect_boxes, img = detect_boxes_rknn(model_path, image_path)
+    angles = get_angles(detect_boxes)
+    for i, angle in enumerate(angles):
+        print(f"框 {i+1}: angle = {angle:.4f} rad ({np.degrees(angle):.2f}°)")
+
+    save_path_all = "./inference_results/boxes_all.jpg"
+    draw_boxes(img.copy(), detect_boxes, save_path_all)
+
+    save_path_top = "./inference_results/top_box.jpg"
+    visualize_top_box(img.copy(), detect_boxes, save_path_top)

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']}")

yemian/resize/rtest.py

@@ -7,7 +7,7 @@ from ultralytics import YOLO
 from pathlib import Path
 
 # ====================== 配置参数 ======================
-MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg_r/exp/weights/best.pt"
+MODEL_PATH = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/seg_r/exp2/weights/best.pt"
 #SOURCE_IMG_DIR = "/home/hx/yolo/output_masks"  # 原始输入图像目录
 SOURCE_IMG_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/f6"  # 原始输入图像目录
 OUTPUT_DIR = "/home/hx/yolo/output_masks2"  # 推理输出根目录

zhuangtai_class_cls/remain_tuili.py

@@ -0,0 +1,143 @@
+import os
+from pathlib import Path
+import cv2
+import numpy as np
+from ultralytics import YOLO
+
+# ---------------------------
+# 类别映射
+# ---------------------------
+CLASS_NAMES = {
+    0: "未堆料",
+    1: "小堆料",
+    2: "大堆料",
+    3: "未浇筑满",
+    4: "浇筑满"
+}
+
+# ---------------------------
+# 加载 ROI 列表
+# ---------------------------
+def load_global_rois(txt_path):
+    rois = []
+    if not os.path.exists(txt_path):
+        print(f"❌ ROI 文件不存在: {txt_path}")
+        return rois
+    with open(txt_path, 'r') as f:
+        for line in f:
+            s = line.strip()
+            if s:
+                try:
+                    x, y, w, h = map(int, s.split(','))
+                    rois.append((x, y, w, h))
+                except Exception as e:
+                    print(f"⚠️ 无法解析 ROI 行 '{s}': {e}")
+    return rois
+
+# ---------------------------
+# 裁剪并 resize ROI
+# ---------------------------
+def crop_and_resize(img, rois, target_size=640):
+    crops = []
+    h_img, w_img = img.shape[:2]
+    for i, (x, y, w, h) in enumerate(rois):
+        if x < 0 or y < 0 or x + w > w_img or y + h > h_img:
+            continue
+        roi = img[y:y+h, x:x+w]
+        roi_resized = cv2.resize(roi, (target_size, target_size), interpolation=cv2.INTER_AREA)
+        crops.append((roi_resized, i))
+    return crops
+
+# ---------------------------
+# class1/class2 加权判断
+# ---------------------------
+def weighted_small_large(pred_probs, threshold=0.4, w1=0.3, w2=0.7):
+    p1 = float(pred_probs[1])
+    p2 = float(pred_probs[2])
+    total = p1 + p2
+    if total > 0:
+        score = (w1 * p1 + w2 * p2) / total
+    else:
+        score = 0.0
+    final_class = "大堆料" if score >= threshold else "小堆料"
+    return final_class, score, p1, p2
+
+# ---------------------------
+# 单张图片推理函数
+# ---------------------------
+def classify_image_weighted(image, model, threshold=0.5):
+    results = model(image)
+    pred_probs = results[0].probs.data.cpu().numpy().flatten()
+    class_id = int(pred_probs.argmax())
+    confidence = float(pred_probs[class_id])
+    class_name = CLASS_NAMES.get(class_id, f"未知类别({class_id})")
+
+    # class1/class2 使用加权得分
+    if class_id in [1, 2]:
+        final_class, score, p1, p2 = weighted_small_large(pred_probs, threshold=threshold)
+    else:
+        final_class = class_name
+        score = confidence
+        p1 = float(pred_probs[1])
+        p2 = float(pred_probs[2])
+
+    return final_class, score, p1, p2
+
+# ---------------------------
+# 批量推理主函数
+# ---------------------------
+def batch_classify_images(model_path, input_folder, output_root, roi_file, target_size=640, threshold=0.5):
+    # 加载模型
+    model = YOLO(model_path)
+
+    # 确保输出根目录存在
+    output_root = Path(output_root)
+    output_root.mkdir(parents=True, exist_ok=True)
+
+    # 为所有类别创建目录
+    class_dirs = {}
+    for name in CLASS_NAMES.values():
+        d = output_root / name
+        d.mkdir(exist_ok=True)
+        class_dirs[name] = d
+
+    rois = load_global_rois(roi_file)
+    if not rois:
+        print("❌ 没有有效 ROI，退出")
+        return
+
+    # 遍历图片
+    for img_path in Path(input_folder).glob("*.*"):
+        if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.bmp', '.tif']:
+            continue
+        try:
+            img = cv2.imread(str(img_path))
+            if img is None:
+                continue
+
+            crops = crop_and_resize(img, rois, target_size)
+
+            for roi_resized, roi_idx in crops:
+                final_class, score, p1, p2 = classify_image_weighted(roi_resized, model, threshold=threshold)
+
+                # 文件名中保存 ROI、类别、加权分数、class1/class2 置信度
+                suffix = f"_roi{roi_idx}_{final_class}_score{score:.2f}_p1{p1:.2f}_p2{p2:.2f}"
+                dst_path = class_dirs[final_class] / f"{img_path.stem}{suffix}{img_path.suffix}"
+                cv2.imwrite(dst_path, roi_resized)
+                print(f"{img_path.name}{suffix} -> {final_class} (score={score:.2f}, p1={p1:.2f}, p2={p2:.2f})")
+
+        except Exception as e:
+            print(f"⚠️ 处理失败 {img_path.name}: {e}")
+
+# ---------------------------
+# 使用示例
+# ---------------------------
+if __name__ == "__main__":
+    model_path = "/home/hx/yolo/ultralytics_yolo11-main/runs/train/cls_resize/exp_cls2/weights/best.pt"
+    input_folder = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/f6"
+    output_root = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/classified"
+    roi_file = "./roi_coordinates/1_rois.txt"
+    target_size = 640
+    threshold = 0.4  # 可调节的比例系数
+
+    batch_classify_images(model_path, input_folder, output_root, roi_file, target_size, threshold)