188 lines
5.8 KiB
Python
Executable File
188 lines
5.8 KiB
Python
Executable File
"""
|
||
快速测试脚本 - 自动查找并可视化第一个有效样本
|
||
|
||
用途：自动从验证集中查找一个带标签的样本并进行可视化，用于测试脚本功能
|
||
|
||
使用方法：
|
||
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()
|