57 lines
2.3 KiB
Python
57 lines
2.3 KiB
Python
import cv2
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from val_labels import load_yolo_polygon_labels,draw_mask_on_image
|
|
|
|
# 之前的load_yolo_polygon_labels和draw_mask_on_image函数保持不变
|
|
|
|
def visualize_labels_in_folder(src_img_dir, src_label_dir, dst_dir, alpha=0.5):
|
|
"""
|
|
遍历源文件夹中的所有图片和对应标签，并在图片上绘制标签后保存至目标文件夹。
|
|
:param src_img_dir: 源图像文件夹路径
|
|
:param src_label_dir: 源标签文件夹路径
|
|
:param dst_dir: 目标文件夹路径，用于保存可视化结果
|
|
:param alpha: 叠加mask的透明度
|
|
"""
|
|
# 确保目标文件夹存在
|
|
Path(dst_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
# 获取所有的图片文件
|
|
img_paths = list(Path(src_img_dir).glob('*.jpg')) + list(Path(src_img_dir).glob('*.png'))
|
|
|
|
for img_path in img_paths:
|
|
txt_path = Path(src_label_dir) / f"{img_path.stem}.txt"
|
|
if not txt_path.exists():
|
|
print(f"⚠️ 跳过 {img_path.name}，未找到对应的标签文件")
|
|
continue
|
|
|
|
# 读取图像
|
|
img = cv2.imread(str(img_path))
|
|
if img is None:
|
|
print(f"❌ 无法读取图像: {img_path}")
|
|
continue
|
|
|
|
h, w = img.shape[:2]
|
|
print(f"🖼️ 正在处理图像: {img_path.name}, 大小: {w}x{h}")
|
|
|
|
# 加载标签
|
|
labels = load_yolo_polygon_labels(txt_path, img.shape)
|
|
if len(labels) == 0:
|
|
print(f"🟡 未找到有效标签: {txt_path}")
|
|
continue
|
|
|
|
# 绘制叠加图
|
|
result_img = draw_mask_on_image(img, labels, alpha)
|
|
|
|
# 保存结果
|
|
output_path = Path(dst_dir) / img_path.name
|
|
cv2.imwrite(str(output_path), result_img)
|
|
print(f"✅ 已保存可视化结果: {output_path}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
SRC_IMG_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize1" # 修改为你的源图像文件夹路径
|
|
SRC_LABEL_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize1" # 修改为你的源标签文件夹路径
|
|
DST_DIR = "/media/hx/04e879fa-d697-4b02-ac7e-a4148876ebb0/dataset/seg/resize_seg5" # 修改为目标文件夹路径
|
|
|
|
visualize_labels_in_folder(SRC_IMG_DIR, SRC_LABEL_DIR, DST_DIR) |