bushu

angle_base_obb/label_view.py

@@ -0,0 +1,120 @@
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+
+# =========================
+# 强制使用非 GUI 后端（关键！）
+# =========================
+import matplotlib
+
+matplotlib.use('Agg')  # 必须在 import pyplot 之前设置
+
+
+def visualize_obb(image_path, label_path, output_dir="output_visualizations"):
+    """
+    可视化图片及其 OBB 标签，并保存结果图像到指定目录。
+
+    :param image_path: 图片路径
+    :param label_path: 标签路径
+    :param output_dir: 输出目录（自动创建）
+    """
+    # 读取图像
+    image = cv2.imread(image_path)
+    if image is None:
+        print(f"❌ 无法读取图像: {image_path}")
+        return
+
+    h, w = image.shape[:2]
+    print(f"✅ 正在处理图像: {os.path.basename(image_path)} | 尺寸: {w} x {h}")
+
+    # 创建用于绘图的副本（BGR → 绘图用）
+    img_draw = image.copy()
+
+    # 读取标签
+    try:
+        with open(label_path, 'r') as f:
+            lines = f.readlines()
+    except Exception as e:
+        print(f"❌ 无法读取标签文件 {label_path}: {e}")
+        return
+
+    for line in lines:
+        parts = line.strip().split()
+        if len(parts) < 9:
+            print(f"⚠️ 跳过无效标签行: {line}")
+            continue
+
+        # 解析：class_id x1 y1 x2 y2 x3 y3 x4 y4
+        try:
+            points = np.array([float(x) for x in parts[1:9]]).reshape(4, 2)
+        except:
+            print(f"⚠️ 坐标解析失败: {line}")
+            continue
+
+        # 归一化坐标 → 像素坐标
+        points[:, 0] *= w  # x
+        points[:, 1] *= h  # y
+        points = np.int32(points)
+
+        # 绘制四边形（绿色）
+        cv2.polylines(img_draw, [points], isClosed=True, color=(0, 255, 0), thickness=3)
+
+        # 绘制顶点（红色圆圈）
+        for (x, y) in points:
+            cv2.circle(img_draw, (x, y), 6, (0, 0, 255), -1)  # 红色实心圆
+
+    # 转为 RGB 用于 matplotlib 保存
+    img_rgb = cv2.cvtColor(img_draw, cv2.COLOR_BGR2RGB)
+
+    # 创建输出目录
+    os.makedirs(output_dir, exist_ok=True)
+
+    # 生成输出路径
+    filename = os.path.splitext(os.path.basename(image_path))[0] + "_vis.png"
+    output_path = os.path.join(output_dir, filename)
+
+    # 使用 matplotlib 保存图像（不显示）
+    plt.figure(figsize=(16, 9), dpi=100)
+    plt.imshow(img_rgb)
+    plt.title(f"OBB Visualization - {os.path.basename(image_path)}", fontsize=14)
+    plt.axis('off')
+    plt.tight_layout()
+    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
+    plt.close()  # 释放内存
+
+    print(f"✅ 可视化结果已保存: {output_path}")
+
+
+def process_directory(directory):
+    """
+    遍历目录，处理所有图片和对应的 .txt 标签文件
+    """
+    print(f"🔍 开始处理目录: {directory}")
+    count = 0
+
+    for filename in os.listdir(directory):
+        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
+            image_path = os.path.join(directory, filename)
+            label_path = os.path.splitext(image_path)[0] + ".txt"
+
+            if os.path.exists(label_path):
+                visualize_obb(image_path, label_path)
+                count += 1
+            else:
+                print(f"🟡 跳过 (无标签): {filename}")
+
+    print(f"🎉 处理完成！共处理 {count} 张图像。")
+
+
+# =========================
+# 主程序入口
+# =========================
+if __name__ == "__main__":
+    # 设置你的数据目录
+    directory = '/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/obb4/labels'
+
+    if not os.path.exists(directory):
+        raise FileNotFoundError(f"目录不存在: {directory}")
+
+    process_directory(directory)