188 lines
5.8 KiB
Python
188 lines
5.8 KiB
Python
|
"""
|
|||
|
|
快速测试脚本 - 自动查找并可视化第一个有效样本
|
|||
|
|
|
|||
|
|
用途：自动从验证集中查找一个带标签的样本并进行可视化，用于测试脚本功能
|
|||
|
|
|
|||
|
|
使用方法：
|
|||
|
|
python scripts_for_gt/test_visualization.py --data data/mono3d.yaml
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
import argparse
|
|||
|
|
import sys
|
|||
|
|
from pathlib import Path
|
|||
|
|
import yaml
|
|||
|
|
|
|||
|
|
# Add project root to path
|
|||
|
|
FILE = Path(__file__).resolve()
|
|||
|
|
ROOT = FILE.parents[1]
|
|||
|
|
if str(ROOT) not in sys.path:
|
|||
|
|
sys.path.append(str(ROOT))
|
|||
|
|
|
|||
|
|
from visualize_single_frame import visualize_single_frame
|
|||
|
|
|
|||
|
|
|
|||
|
|
def find_first_valid_sample(data_yaml_path):
|
|||
|
|
"""从数据集配置中查找第一个有效样本
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
data_yaml_path: 数据集YAML配置文件路径
|
|||
|
|
|
|||
|
|
Returns:
|
|||
|
|
tuple: (image_path, label_path, calib_path) 或 None
|
|||
|
|
"""
|
|||
|
|
# 读取数据集配置
|
|||
|
|
with open(data_yaml_path, 'r') as f:
|
|||
|
|
data = yaml.safe_load(f)
|
|||
|
|
|
|||
|
|
# 获取验证集路径
|
|||
|
|
val_path = data.get('val')
|
|||
|
|
if val_path is None:
|
|||
|
|
print("错误：数据集配置中未找到'val'字段")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
# 如果是相对路径，相对于yaml文件的位置
|
|||
|
|
val_path = Path(val_path)
|
|||
|
|
if not val_path.is_absolute():
|
|||
|
|
val_path = Path(data_yaml_path).parent / val_path
|
|||
|
|
|
|||
|
|
# 读取验证集图像列表
|
|||
|
|
if val_path.is_file():
|
|||
|
|
# 如果是文件列表
|
|||
|
|
with open(val_path, 'r') as f:
|
|||
|
|
image_files = [line.strip() for line in f.readlines() if line.strip()]
|
|||
|
|
else:
|
|||
|
|
# 如果是目录
|
|||
|
|
image_files = list(val_path.glob("**/*.jpg"))
|
|||
|
|
image_files.extend(list(val_path.glob("**/*.png")))
|
|||
|
|
|
|||
|
|
print(f"在验证集中找到 {len(image_files)} 个图像文件")
|
|||
|
|
|
|||
|
|
# 遍历查找第一个有标签的样本
|
|||
|
|
for image_path in image_files:
|
|||
|
|
image_rel_path = Path(image_path)
|
|||
|
|
|
|||
|
|
image_path = val_path.parent / image_rel_path if not image_rel_path.is_absolute() else image_rel_path
|
|||
|
|
# 推断标签路径
|
|||
|
|
label_path = Path(str(image_path).replace('/images/', '/labels/').replace('.jpg', '.txt').replace('.png', '.txt'))
|
|||
|
|
|
|||
|
|
# 推断标定路径
|
|||
|
|
calib_path = image_path.parent.parent / 'calib' / 'L2_calib' / 'camera4.json'
|
|||
|
|
|
|||
|
|
if label_path.exists():
|
|||
|
|
print(f"\n找到有效样本:")
|
|||
|
|
print(f" 图像: {image_path}")
|
|||
|
|
print(f" 标签: {label_path}")
|
|||
|
|
print(f" 标定: {calib_path if calib_path.exists() else '未找到'}")
|
|||
|
|
return str(image_path), str(label_path), str(calib_path) if calib_path.exists() else None
|
|||
|
|
|
|||
|
|
print("错误：未找到有效样本")
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def test_visualization(data_yaml_path, output_dir="./test_gt_viz", use_roi=True):
|
|||
|
|
"""测试可视化功能
|
|||
|
|
|
|||
|
|
Args:
|
|||
|
|
data_yaml_path: 数据集YAML配置文件路径
|
|||
|
|
output_dir: 输出目录
|
|||
|
|
use_roi: 是否使用ROI变换
|
|||
|
|
"""
|
|||
|
|
# 读取数据集配置
|
|||
|
|
with open(data_yaml_path, 'r') as f:
|
|||
|
|
data = yaml.safe_load(f)
|
|||
|
|
|
|||
|
|
# 查找第一个有效样本
|
|||
|
|
result = find_first_valid_sample(data_yaml_path)
|
|||
|
|
if result is None:
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
image_path, label_path, calib_path = result
|
|||
|
|
|
|||
|
|
# 获取配置参数
|
|||
|
|
roi = data.get('roi') # [width, height]
|
|||
|
|
virtual_fx = data.get('virtual_fx')
|
|||
|
|
ori_img_size = data.get('ori_img_size') # [width, height]
|
|||
|
|
|
|||
|
|
print(f"\n数据集配置:")
|
|||
|
|
print(f" ROI: {roi}")
|
|||
|
|
print(f" 虚拟焦距: {virtual_fx}")
|
|||
|
|
print(f" 原始图像尺寸: {ori_img_size}")
|
|||
|
|
|
|||
|
|
# 类别名称
|
|||
|
|
names = {
|
|||
|
|
0: "vehicle",
|
|||
|
|
1: "pedestrian",
|
|||
|
|
2: "bicycle",
|
|||
|
|
3: "rider",
|
|||
|
|
13: "tricycle",
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
print(f"\n开始可视化...")
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
# 测试1：不使用ROI
|
|||
|
|
print("\n=== 测试1：原始图像可视化（不使用ROI）===")
|
|||
|
|
visualize_single_frame(
|
|||
|
|
image_path=image_path,
|
|||
|
|
label_path=label_path,
|
|||
|
|
output_dir=f"{output_dir}/no_roi",
|
|||
|
|
calib_path=calib_path,
|
|||
|
|
roi_size=None,
|
|||
|
|
virtual_fx=None,
|
|||
|
|
ori_img_size=None,
|
|||
|
|
names=names,
|
|||
|
|
)
|
|||
|
|
print("✓ 测试1完成")
|
|||
|
|
|
|||
|
|
# 测试2：使用ROI（如果配置了）
|
|||
|
|
if use_roi and roi is not None and virtual_fx is not None:
|
|||
|
|
print("\n=== 测试2：ROI变换后可视化 ===")
|
|||
|
|
visualize_single_frame(
|
|||
|
|
image_path=image_path,
|
|||
|
|
label_path=label_path,
|
|||
|
|
output_dir=f"{output_dir}/with_roi",
|
|||
|
|
calib_path=calib_path,
|
|||
|
|
roi_size=tuple(roi),
|
|||
|
|
virtual_fx=virtual_fx,
|
|||
|
|
ori_img_size=tuple(ori_img_size) if ori_img_size else None,
|
|||
|
|
names=names,
|
|||
|
|
)
|
|||
|
|
print("✓ 测试2完成")
|
|||
|
|
|
|||
|
|
print(f"\n✅ 所有测试完成！结果保存在: {output_dir}")
|
|||
|
|
print("\n生成的文件:")
|
|||
|
|
output_path = Path(output_dir)
|
|||
|
|
for file in sorted(output_path.rglob("*.jpg")):
|
|||
|
|
print(f" {file.relative_to(output_path)}")
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"\n❌ 测试失败: {e}")
|
|||
|
|
import traceback
|
|||
|
|
traceback.print_exc()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def parse_args():
|
|||
|
|
parser = argparse.ArgumentParser(description="快速测试真值可视化功能")
|
|||
|
|
parser.add_argument("--data", type=str, default="data/mono3d.yaml",
|
|||
|
|
help="数据集YAML配置文件")
|
|||
|
|
parser.add_argument("--output", type=str, default="./test_gt_viz",
|
|||
|
|
help="输出目录")
|
|||
|
|
parser.add_argument("--no-roi", action="store_true",
|
|||
|
|
help="不测试ROI变换")
|
|||
|
|
|
|||
|
|
return parser.parse_args()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def main():
|
|||
|
|
args = parse_args()
|
|||
|
|
|
|||
|
|
test_visualization(
|
|||
|
|
data_yaml_path=args.data,
|
|||
|
|
output_dir=args.output,
|
|||
|
|
use_roi=not args.no_roi,
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
main()
|