v0.1.17; update 7-9

robogauge/utils/visualize/plot_terrain_levels.py

@@ -0,0 +1,168 @@
+# -*- coding: utf-8 -*-
+'''
+@File    : plot_terrain_levels.py
+@Time    : 2025/12/27 23:04:01
+@Author  : wty-yy (with Gemini 3)
+@Version : 1.0
+@Blog    : https://wty-yy.github.io/
+@Desc    : None
+@Desc    : 可视化不同地形下随着摩擦力变化的关卡通过等级 (Terrain Level vs Friction)
+'''
+import matplotlib.pyplot as plt
+from pathlib import Path
+import numpy as np
+import yaml
+import argparse
+import os
+import re
+
+# --- 配置 Matplotlib 样式 ---
+config = {
+    "font.family": 'serif',
+    "figure.figsize": (18, 10),
+    "font.size": 12,
+    "mathtext.fontset": 'cm',
+    'axes.unicode_minus': False
+}
+plt.rcParams.update(config)
+plt.rcParams['font.serif'] = ['Times New Roman', 'DejaVu Serif', 'serif']
+
+def load_terrain_data(file_paths):
+    """
+    加载数据并按模型、地形、摩擦力组织
+    Returns:
+        model_data: {model_name: {terrain_type: {friction: level}}}
+    """
+    target_terrains = ['wave', 'slope', 'stairs_up', 'stairs_down', 'obstacle']
+    model_data = {}
+    
+    # 用于从键名中提取 friction 的正则 (例如 friction1.25)
+    friction_pattern = re.compile(r'friction([\d\.]+)')
+
+    for path in file_paths:
+        if not os.path.exists(path):
+            print(f"[ERROR] File not found: {path}")
+            continue
+            
+        with open(path, 'r', encoding='utf-8') as f:
+            content = yaml.safe_load(f)
+        
+        # 获取模型名称
+        raw_path = content.get('model_path', 'Unknown_Model')
+        model_name = os.path.basename(raw_path).replace('.pt', '')
+        
+        if model_name not in model_data:
+            model_data[model_name] = {t: {} for t in target_terrains}
+            
+        # 遍历YAML中的每个测试条目
+        for key, value in content.items():
+            if key in ['model_path', 'summary']:
+                continue
+            
+            # 1. 检查地形是否在目标列表中
+            t_name = value.get('terrain_name')
+            if t_name not in target_terrains:
+                continue
+            
+            # 2. 从 key 中提取摩擦力
+            match = friction_pattern.search(key)
+            if not match:
+                continue
+            friction = float(match.group(1))
+            
+            # 3. 获取地形等级 (取该条件下该地形的最大值，防止重复)
+            level = value.get('terrain_level', 0)
+            
+            # 记录数据
+            current_max = model_data[model_name][t_name].get(friction, -1)
+            if level > current_max:
+                model_data[model_name][t_name][friction] = level
+                
+    return model_data, target_terrains
+
+def plot_results(model_data, terrain_labels, output_file=None):
+    """绘制 2x3 的子图"""
+    if not model_data:
+        print("No data to plot.")
+        return
+
+    # 设置 2行3列 的布局
+    fig, axes = plt.subplots(2, 3, figsize=(18, 10), sharex=True, sharey=True)
+    axes = axes.flatten()
+    
+    # 自动获取颜色映射
+    colors = plt.cm.get_cmap("tab10", len(model_data))
+    
+    # 遍历前5个地形进行绘图
+    for i, t_name in enumerate(terrain_labels):
+        ax = axes[i]
+        
+        for idx, (model_name, terrains) in enumerate(model_data.items()):
+            data_points = terrains.get(t_name, {})
+            if not data_points:
+                continue
+            
+            # 排序 (Friction, Level)
+            sorted_points = sorted(data_points.items())
+            xs, ys = zip(*sorted_points)
+            
+            ax.plot(xs, ys, marker='o', linestyle='-', linewidth=2, 
+                    label=model_name, color=colors(idx), alpha=0.8)
+        
+        ax.set_title(t_name.replace('_', ' ').title(), fontsize=14)
+        ax.grid(True, linestyle='--', alpha=0.6)
+
+    # --- 第6张图：绘制平均值 (Average) ---
+    ax_avg = axes[5]
+    for idx, (model_name, terrains) in enumerate(model_data.items()):
+        # 聚合该模型下所有地形的数据用于计算平均值
+        # 结构: friction -> [level_slope, level_wave, ...]
+        friction_agg = {}
+        for t_name in terrain_labels:
+            for f, l in terrains[t_name].items():
+                if f not in friction_agg:
+                    friction_agg[f] = []
+                friction_agg[f].append(l)
+        
+        if not friction_agg:
+            continue
+            
+        # 计算平均值并排序
+        sorted_frictions = sorted(friction_agg.keys())
+        avg_levels = [np.mean(friction_agg[f]) for f in sorted_frictions]
+        
+        ax_avg.plot(sorted_frictions, avg_levels, marker='s', linestyle='-', linewidth=2,
+                    label=model_name, color=colors(idx), alpha=0.8)
+
+    ax_avg.set_title('Average Performance', fontsize=14)
+    ax_avg.grid(True, linestyle='--', alpha=0.6)
+    
+    # 仅在最后一张图显示图例，避免遮挡
+    ax_avg.legend(loc='best', frameon=False, title="Models")
+
+    # 最下面一行设置 X 轴标签
+    for i in [3, 4, 5]:
+        axes[i].set_xlabel('Friction ($\mu$)')
+
+    # 最左边一列设置 Y 轴标签
+    for i in [0, 3]:
+        axes[i].set_ylabel('Terrain Level')
+
+    plt.tight_layout()
+    # 稍微调整间距防止标签重叠
+    # plt.subplots_adjust(wspace=0.1, hspace=0.15)
+    
+    if output_file:
+        plt.savefig(output_file, dpi=300, bbox_inches='tight')
+        print(f"Plot saved to {output_file}")
+    else:
+        plt.show()
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Plot terrain level vs friction curves.")
+    parser.add_argument('files', metavar='F', type=str, nargs='+', help='YAML result files')
+    parser.add_argument('--out', type=str, default='terrain_analysis.jpg', help='Output image filename')
+    args = parser.parse_args()
+    
+    data, labels = load_terrain_data(args.files)
+    plot_results(data, labels, output_file=args.out)