125 lines
4.0 KiB
Python
125 lines
4.0 KiB
Python
# debug_label_overlay.py
|
||
import cv2
|
||
import numpy as np
|
||
from pathlib import Path
|
||
|
||
def load_yolo_polygon_labels(txt_path, img_shape):
|
||
"""
|
||
加载 YOLO 格式的多边形标签 (class x1 y1 x2 y2 ... xn yn)
|
||
:param txt_path: .txt 标签路径
|
||
:param img_shape: (H, W) 原图大小
|
||
:return: 列表，元素为 [cls_id, polygon]，polygon 是 Nx2 的点阵
|
||
"""
|
||
h, w = img_shape[:2]
|
||
labels = []
|
||
if not Path(txt_path).exists():
|
||
print(f"⚠️ 标签文件不存在: {txt_path}")
|
||
return labels
|
||
|
||
with open(txt_path, 'r') as f:
|
||
for line in f:
|
||
line = line.strip()
|
||
if not line or line.startswith('#'):
|
||
continue
|
||
parts = list(map(float, line.split()))
|
||
cls_id = int(parts[0])
|
||
coords = np.array(parts[1:]).reshape(-1, 2) # (N, 2)
|
||
# 反归一化到原图
|
||
coords[:, 0] *= w
|
||
coords[:, 1] *= h
|
||
coords = coords.astype(np.int32)
|
||
labels.append((cls_id, coords))
|
||
return labels
|
||
|
||
def draw_mask_on_image(img, labels, alpha=0.5):
|
||
"""
|
||
在图像上绘制半透明 mask 和边框
|
||
"""
|
||
overlay = img.copy()
|
||
colors = [
|
||
(0, 255, 0), # 绿
|
||
(0, 0, 255), # 蓝
|
||
(255, 0, 0), # 红
|
||
(255, 255, 0),
|
||
(255, 0, 255),
|
||
(0, 255, 255),
|
||
(128, 0, 128),
|
||
(0, 128, 0),
|
||
(0, 0, 128),
|
||
(128, 128, 0)
|
||
]
|
||
|
||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||
font_scale = 0.7
|
||
thickness = 1
|
||
|
||
for cls_id, poly in labels:
|
||
color = colors[cls_id % len(colors)]
|
||
|
||
# 填充半透明 mask
|
||
cv2.fillPoly(overlay, [poly], color=color)
|
||
|
||
# 绘制边缘轮廓
|
||
cv2.polylines(overlay, [poly], isClosed=True, color=color, thickness=2)
|
||
|
||
# 添加类别标签
|
||
x, y = poly[0][0], poly[0][1] # 使用第一个点作为标签位置
|
||
label = f"Class {cls_id}"
|
||
(text_w, text_h), _ = cv2.getTextSize(label, font, font_scale, thickness)
|
||
cv2.rectangle(overlay, (x, y - text_h - 6), (x + text_w, y), color, -1)
|
||
cv2.putText(overlay, label, (x, y - 4), font, font_scale,
|
||
(255, 255, 255), thickness, cv2.LINE_AA)
|
||
|
||
# 混合叠加
|
||
result = cv2.addWeighted(overlay, alpha, img, 1 - alpha, 0)
|
||
return result
|
||
|
||
def main():
|
||
# ====================== 配置区 ======================
|
||
IMG_PATH = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/yemian_seg_camera01/train/frame_20250805_121717_22320.jpg" # 修改为你的某张测试图
|
||
TXT_PATH = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/yemian_seg_camera01/train/frame_20250805_121717_22320.txt" # 对应的 .txt 文件
|
||
OUTPUT_PATH = "/home/hx/yolo/output_masks2/debug_overlay.png"
|
||
# ====================================================
|
||
|
||
img_path = Path(IMG_PATH)
|
||
txt_path = Path(TXT_PATH)
|
||
output_path = Path(OUTPUT_PATH)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
||
if not img_path.exists():
|
||
print(f"❌ 图像文件不存在: {img_path}")
|
||
return
|
||
|
||
# 读取图像
|
||
img = cv2.imread(str(img_path))
|
||
if img is None:
|
||
print(f"❌ 无法读取图像: {img_path}")
|
||
return
|
||
|
||
h, w = img.shape[:2]
|
||
print(f"🖼️ 图像大小: {w}x{h}")
|
||
|
||
# 加载标签
|
||
labels = load_yolo_polygon_labels(txt_path, img.shape)
|
||
if len(labels) == 0:
|
||
print(f"🟡 未找到有效标签: {txt_path}")
|
||
return
|
||
|
||
print(f"🎯 加载到 {len(labels)} 个对象")
|
||
|
||
# 绘制叠加图
|
||
result_img = draw_mask_on_image(img, labels, alpha=0.5)
|
||
|
||
# 保存结果
|
||
cv2.imwrite(str(output_path), result_img)
|
||
print(f"✅ 已保存可视化结果: {output_path}")
|
||
|
||
# 显示（可选）
|
||
display_img = cv2.resize(result_img, (960, 540)) # 缩放便于显示
|
||
cv2.imshow("Label Overlay Debug", display_img)
|
||
print("📌 按任意键关闭窗口...")
|
||
cv2.waitKey(0)
|
||
cv2.destroyAllWindows()
|
||
|
||
if __name__ == "__main__":
|
||
main() |