增加cvat反向上传

tool/save_bigangle.py

@@ -0,0 +1,117 @@
+import cv2
+import os
+import shutil
+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,
+    large_angle_output_dir="./large_angle_images",  # 新增：大角度图像保存目录
+    conf_thresh=0.15,
+    imgsz=640,
+    angle_threshold_deg=10.0  # 夹角阈值：超过即视为“大角度”
+):
+    """
+    批量处理 OBB 图像，若任意两个目标的主方向夹角 > angle_threshold_deg，
+    则将原图移动到 large_angle_output_dir。
+    """
+    os.makedirs(large_angle_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]
+
+        # 提取旋转角
+        boxes = result.obb
+        angles_deg = []
+        has_large_angle = False  # 标记是否有大角度的目标
+        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)
+                        angle_diff_deg = np.degrees(min_diff_rad)
+                        pairwise_angles_deg.append(angle_diff_deg)
+                        print(f"  Box {i + 1} 与 Box {j + 1} 夹角: {angle_diff_deg:.2f}°")
+                        if angle_diff_deg > angle_threshold_deg:
+                            has_large_angle = True
+            elif len(angles_deg) == 1:
+                print("  仅检测到一个目标，无法计算夹角")
+
+        # 保存结果
+        results_dict[img_filename] = {
+            "angles_deg": angles_deg,
+            "pairwise_angles_deg": pairwise_angles_deg,
+            "has_large_angle": has_large_angle
+        }
+
+        # 如果存在大角度，移动原图到新文件夹
+        if has_large_angle:
+            print(f"发现夹角 > {angle_threshold_deg}°，移动原图到大角度文件夹")
+            shutil.move(img_path, os.path.join(large_angle_output_dir, img_filename))
+
+    print(f"\n所有图像处理完成！大角度图像已移动至: {large_angle_output_dir}")
+    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/1/ready"
+    LARGE_ANGLE_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/1/large_angle_images"  # 新增输出目录
+
+    results = process_obb_images(
+        model_path=MODEL_PATH,
+        image_dir=IMAGE_SOURCE_DIR,
+        large_angle_output_dir=LARGE_ANGLE_DIR,
+        conf_thresh=0.15,
+        imgsz=640,
+        angle_threshold_deg=10.0
+    )
+
+    # 可选：打印大角度图像
+    large_angle_imgs = [name for name, info in results.items() if info["has_large_angle"]]
+    print(f"\n共 {len(large_angle_imgs)} 张图像包含 >10° 的夹角：")
+    for name in large_angle_imgs:
+        print(f"  - {name}")