120 lines
3.6 KiB
Python
120 lines
3.6 KiB
Python
|
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)
|