#!/usr/bin/env python3 import json import os import sys import glob import argparse import matplotlib.pyplot as plt import pandas as pd METRIC_TITLE_MAPPING = { 'docvqa_val_anls': 'DocVQA', 'infovqa_val_anls': 'InfoVQA', 'mme_total_score': 'MME Total', 'mmmu_val_mmmu_acc': 'MMMU', 'mmstar_average': 'MMStar', 'ocrbench_ocrbench_accuracy': 'OCRBench', 'scienceqa_exact_match': 'ScienceQA', 'textvqa_val_exact_match': 'TextVQA', 'average': 'Average', 'average_rank': 'Average Rank', 'ai2d_exact_match': 'AI2D', 'chartqa_relaxed_overall': 'ChartQA', 'seedbench_seed_all': 'SeedBench' } def compute_ranking_summary(all_results, tasks_to_plot): """Compute ranking-based summary metric across all runs.""" if not all_results or len(all_results) < 2: return all_results # Get all steps that appear in all runs all_steps = set() for results in all_results: all_steps.update(result['step'] for result in results) # For each step, compute rankings for step in all_steps: # Find all runs that have this step step_data = [] run_indices = [] for run_idx, results in enumerate(all_results): step_result = next((r for r in results if r['step'] == step), None) if step_result: step_data.append(step_result) run_indices.append(run_idx) if len(step_data) < 2: continue # Get metrics to rank (exclude 'average' and 'average_rank' from ranking calculation) metrics_to_rank = [] if tasks_to_plot: for task in tasks_to_plot: if task not in ['average', 'average_rank'] and task in step_data[0]: metrics_to_rank.append(task) else: metrics_to_rank = [k for k in step_data[0].keys() if k not in ['step', 'average', 'average_rank']] if not metrics_to_rank: continue # Compute rankings for each metric rankings = [] for metric in metrics_to_rank: # Get values for this metric across all runs at this step metric_values = [] for data in step_data: if metric in data and isinstance(data[metric], (int, float)): metric_values.append(data[metric]) else: metric_values.append(None) # Skip this metric if any run is missing it if None in metric_values: continue # Create ranking (higher value = better rank, so we rank in descending order) # Convert to list of (value, original_index) pairs indexed_values = [(val, idx) for idx, val in enumerate(metric_values)] # Sort by value in descending order (higher is better) indexed_values.sort(key=lambda x: x[0], reverse=True) # Assign ranks (1 is best) metric_rankings = [0] * len(metric_values) for rank, (_, original_idx) in enumerate(indexed_values, 1): metric_rankings[original_idx] = rank rankings.append(metric_rankings) # Compute average ranking for each run if rankings: avg_rankings = [] for run_idx in range(len(step_data)): run_ranks = [ranking[run_idx] for ranking in rankings] avg_rankings.append(sum(run_ranks) / len(run_ranks)) # Add ranking summary to each run's data for i, (data, run_idx) in enumerate(zip(step_data, run_indices)): # Find the result in the original data and add ranking summary for result in all_results[run_idx]: if result['step'] == step: result['average_rank'] = avg_rankings[i] break return all_results def load_eval_results(eval_folder, tasks_to_plot=None): """Load all JSON files from the evaluation folder and extract results.""" json_files = glob.glob(os.path.join(eval_folder, "step_*.json")) if not json_files: print(f"No JSON files found in {eval_folder}") return None results = [] for json_file in json_files: with open(json_file, 'r') as f: data = json.load(f) step = data.get('global_step', 0) metrics = data.get('results', {}) result = {'step': step} result.update(metrics) # Add MME total score if mme is in tasks and both perception and cognition scores exist if tasks_to_plot and any('mme_total_score' in task.lower() for task in tasks_to_plot): perception_score = result.get('mme_mme_perception_score') cognition_score = result.get('mme_mme_cognition_score') if perception_score is not None and cognition_score is not None: result['mme_total_score'] = perception_score + cognition_score # Add average score if 'average' is in tasks if tasks_to_plot and 'average' in tasks_to_plot: # Get only the specified tasks (excluding 'average') metrics_to_average = [] for task in tasks_to_plot: if (task != 'average' and 'rank' not in task.lower() and task in result and isinstance(result[task], (int, float))): # Special handling for MME total score: normalize by dividing by 2800 if task == 'mme_total_score': normalized_score = result[task] / 2800.0 metrics_to_average.append(normalized_score) else: metrics_to_average.append(result[task]) if metrics_to_average: result['average'] = sum(metrics_to_average) / len(metrics_to_average) results.append(result) # Sort by step results.sort(key=lambda x: x['step']) return results def get_legend_name(eval_folder, custom_name=None): """Extract legend name from folder path or use custom name.""" if custom_name: return custom_name folder_name = os.path.basename(eval_folder) return folder_name.split('_')[-1] def plot_results(all_results, eval_folders, custom_names=None, tasks_to_plot=None, output_filename=None, steps_to_plot=None): """Plot the evaluation results for multiple folders.""" if not all_results: return # Set academic style plt.rcParams['font.family'] = 'serif' plt.rcParams['font.size'] = 12 plt.rcParams['mathtext.fontset'] = 'cm' # Mapping from metric names to display titles # Extract all metric names from all results metric_names = set() for results in all_results: for result in results: metric_names.update(k for k in result.keys() if k != 'step') # Filter metrics based on specified tasks if provided if tasks_to_plot: filtered_metrics = set() for task in tasks_to_plot: # Exact match for specified tasks if task in metric_names: filtered_metrics.add(task) metric_names = filtered_metrics if not metric_names: print(f"Warning: No metrics found exactly matching tasks: {tasks_to_plot}") return metric_names = sorted(list(metric_names)) # Create subplots n_metrics = len(metric_names) n_cols = 3 n_rows = (n_metrics + n_cols - 1) // n_cols _, axes = plt.subplots(n_rows, n_cols, figsize=(15, 4*n_rows)) if n_rows == 1: axes = axes.reshape(1, -1) # Define academic colors and markers for different runs colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf'] markers = ['o', 's', '^', 'D', 'v', 'p', '*', 'h', '+', 'x'] for i, metric in enumerate(metric_names): row = i // n_cols col = i % n_cols ax = axes[row, col] # Plot each run for j, (results, eval_folder) in enumerate(zip(all_results, eval_folders)): # Extract values for this metric values = [] metric_steps = [] missing_steps = [] for result in results: # Check if we should include this step if steps_to_plot is None or result['step'] in steps_to_plot: if metric in result: values.append(result[metric]) metric_steps.append(result['step']) elif steps_to_plot is not None: # Only log missing if specific steps were requested missing_steps.append(result['step']) # Log missing metrics for specified steps if missing_steps: folder_name = custom_names[j] if custom_names and custom_names[j] else os.path.basename(eval_folder) print(f"Warning: {folder_name} missing '{metric}' for steps: {missing_steps}") if values: custom_name = custom_names[j] if custom_names else None legend_name = get_legend_name(eval_folder, custom_name) color = colors[j % len(colors)] marker = markers[j % len(markers)] # Check if there's stderr data for this metric stderr_metric = metric + '_stderr' stderr_values = [] for result in results: if steps_to_plot is None or result['step'] in steps_to_plot: if metric in result and stderr_metric in result: stderr_values.append(result[stderr_metric]) elif metric in result: stderr_values.append(0) # No stderr available for this step # Plot the main line ax.plot(metric_steps, values, marker=marker, markersize=4, color=color, label=legend_name, linewidth=2, alpha=0.9) # Plot error corridor if stderr data is available if stderr_values and any(stderr > 0 for stderr in stderr_values): lower_bounds = [v - s for v, s in zip(values, stderr_values)] upper_bounds = [v + s for v, s in zip(values, stderr_values)] ax.fill_between(metric_steps, lower_bounds, upper_bounds, color=color, alpha=0.2, linewidth=0) # Get display title from mapping, fallback to original metric name display_title = METRIC_TITLE_MAPPING.get(metric, metric) ax.set_title(display_title, fontsize=13, weight='bold') ax.set_xlabel('Training Step (×1000)', fontsize=10, weight='bold') ax.set_ylabel('Value', fontsize=11, weight='bold') ax.grid(True, alpha=0.2, linestyle='--', linewidth=0.5) # Set x-axis ticks to show simple integers (steps divided by 1000) # Get all steps for this metric across all runs all_metric_steps = [] for results in all_results: for result in results: if (steps_to_plot is None or result['step'] in steps_to_plot) and metric in result: all_metric_steps.append(result['step']) if all_metric_steps: unique_steps = sorted(set(all_metric_steps)) ax.set_xticks(unique_steps) # Show only every second label to avoid crowding labels = [int(step/1000) if i % 2 == 0 else '' for i, step in enumerate(unique_steps)] ax.set_xticklabels(labels) # Invert y-axis for ranking metrics (lower rank = better performance) if 'rank' in metric.lower(): ax.invert_yaxis() # Add subtle background and improve spines ax.set_facecolor('#fafafa') for spine in ax.spines.values(): spine.set_linewidth(1.2) spine.set_color('#333333') if len(eval_folders) > 1: ax.legend(loc='upper left', frameon=True, fancybox=True, shadow=False, framealpha=0.9, edgecolor='gray', fontsize=10) # Hide unused subplots for i in range(n_metrics, n_rows * n_cols): row = i // n_cols col = i % n_cols axes[row, col].set_visible(False) # # Add title if output filename is specified # if output_filename: # plt.suptitle(output_filename, fontsize=16, y=0.98) plt.tight_layout() # Create assets folder if it doesn't exist assets_folder = '/fsx/nazar_drugov/nanoVLM_encoder_free/plots_final' os.makedirs(assets_folder, exist_ok=True) # Save the plot to assets folder if output_filename: output_file = os.path.join(assets_folder, f'{output_filename}.pdf') elif len(eval_folders) == 1: folder_name = os.path.basename(eval_folders[0]) output_file = os.path.join(assets_folder, f'{folder_name}_evaluation_plots.pdf') else: output_file = os.path.join(assets_folder, 'comparison_evaluation_plots.pdf') plt.savefig(output_file, format='pdf', dpi=600, bbox_inches='tight') print(f"Plot saved to: {output_file}") plt.close() # Save individual plots as PDFs for specified metrics individual_plots = ['average_rank', 'average'] # Add more metrics here as needed for metric in individual_plots: if metric in metric_names: save_individual_plot_pdf(all_results, eval_folders, custom_names, output_filename, metric, steps_to_plot) # Save CSV data save_csv_data(all_results, eval_folders, custom_names, metric_names, output_file, steps_to_plot) def save_individual_plot_pdf(all_results, eval_folders, custom_names, output_filename, metric_name, steps_to_plot=None): """Save an individual metric plot as a PDF with 300 DPI and no title.""" # Set academic style plt.rcParams['font.family'] = 'serif' plt.rcParams['font.size'] = 12 plt.rcParams['mathtext.fontset'] = 'cm' # Create a new figure with golden ratio proportions #plt.figure(figsize=(10, 6.18)) # Define academic colors and markers colors = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf'] markers = ['o', 's', '^', 'D', 'v', 'p', '*', 'h', '+', 'x'] # Plot each run for the specified metric for j, (results, eval_folder) in enumerate(zip(all_results, eval_folders)): # Extract values for this metric values = [] metric_steps = [] for result in results: # Check if we should include this step if steps_to_plot is None or result['step'] in steps_to_plot: if metric_name in result: values.append(result[metric_name]) metric_steps.append(result['step']) if values: custom_name = custom_names[j] if custom_names else None legend_name = get_legend_name(eval_folder, custom_name) color = colors[j % len(colors)] marker = markers[j % len(markers)] # Check if there's stderr data for this metric stderr_metric = metric_name + '_stderr' stderr_values = [] for result in results: if steps_to_plot is None or result['step'] in steps_to_plot: if metric_name in result and stderr_metric in result: stderr_values.append(result[stderr_metric]) elif metric_name in result: stderr_values.append(0) # No stderr available for this step # Plot the main line plt.plot(metric_steps, values, marker=marker, markersize=6, color=color, label=legend_name, linewidth=2.5, alpha=0.9) # Plot error corridor if stderr data is available if stderr_values and any(stderr > 0 for stderr in stderr_values): lower_bounds = [v - s for v, s in zip(values, stderr_values)] upper_bounds = [v + s for v, s in zip(values, stderr_values)] plt.fill_between(metric_steps, lower_bounds, upper_bounds, color=color, alpha=0.2, linewidth=0) # Configure the plot plt.xlabel('Training Step (×1000)', fontsize=13, weight='bold') display_title = METRIC_TITLE_MAPPING.get(metric_name, metric_name) plt.ylabel(display_title, fontsize=13, weight='bold') plt.grid(True, alpha=0.2, linestyle='--', linewidth=0.5) # Set x-axis limits from 1000 to last datapoint with slight margins all_steps = [] for results in all_results: for result in results: # Only include steps that match our filter and have the metric if (steps_to_plot is None or result['step'] in steps_to_plot) and metric_name in result: all_steps.append(result['step']) if all_steps: min_step = 1000 max_step = max(all_steps) x_margin = (max_step - min_step) * 0.02 # 2% margin plt.xlim(min_step - x_margin, max_step + x_margin) # Set x-axis ticks to show simple integers (steps divided by 1000) unique_steps = sorted(set(all_steps)) #unique_steps = [i for i in range(1000, 60000, 4000)] plt.xticks(unique_steps, [int(step/1000) for step in unique_steps]) # Invert y-axis for ranking metrics (lower rank = better performance) if 'rank' in metric_name.lower(): plt.gca().invert_yaxis() # Set y-axis limits from 1 to number of runs with slight margins y_margin = 0.1 plt.ylim(len(eval_folders) + y_margin, 1 - y_margin) # Add legend if multiple runs if len(eval_folders) > 1: plt.legend(loc='upper left', frameon=True, fancybox=True, shadow=False, framealpha=0.9, edgecolor='gray', fontsize=11) # Add subtle background and improve spines ax = plt.gca() ax.set_facecolor('#fafafa') for spine in ax.spines.values(): spine.set_linewidth(1.2) spine.set_color('#333333') plt.tight_layout(pad=1.5) # Create assets folder if it doesn't exist assets_folder = '/fsx/nazar_drugov/nanoVLM_encoder_free/plots_final' os.makedirs(assets_folder, exist_ok=True) # Generate filename for individual plot PDF if output_filename: pdf_file = os.path.join(assets_folder, f'{output_filename}_{metric_name}.pdf') elif len(eval_folders) == 1: folder_name = os.path.basename(eval_folders[0]) pdf_file = os.path.join(assets_folder, f'{folder_name}_{metric_name}.pdf') else: pdf_file = os.path.join(assets_folder, f'comparison_{metric_name}.pdf') # Save as PDF with 300 DPI plt.savefig(pdf_file, format='pdf', dpi=300, bbox_inches='tight') print(f"Individual plot for '{metric_name}' saved to: {pdf_file}") # Also save as PNG png_file = pdf_file.replace('.pdf', '.png') plt.savefig(png_file, format='png', dpi=300, bbox_inches='tight') print(f"Individual plot for '{metric_name}' saved to: {png_file}") plt.close() def save_csv_data(all_results, eval_folders, custom_names, metric_names, output_file, steps_to_plot=None): """Save the plot data to a CSV file.""" # Prepare data for CSV csv_data = [] for i, (results, eval_folder) in enumerate(zip(all_results, eval_folders)): # Get the run name custom_name = custom_names[i] if custom_names else None run_name = get_legend_name(eval_folder, custom_name) for result in results: step = result['step'] # Only include steps that are plotted if steps_to_plot is None or step in steps_to_plot: for metric in metric_names: if metric in result: row_data = { 'run': run_name, 'step': step, 'metric': metric, 'value': result[metric] } # Add stderr if available stderr_metric = metric + '_stderr' if stderr_metric in result: row_data['stderr'] = result[stderr_metric] csv_data.append(row_data) # Convert to DataFrame and save if csv_data: df = pd.DataFrame(csv_data) # Generate CSV filename from plot filename csv_file = output_file.replace('.pdf', '.csv') df.to_csv(csv_file, index=False) print(f"Data saved to: {csv_file}") def parse_args(): """Parse command line arguments supporting both folder and folder:name format.""" parser = argparse.ArgumentParser( description='Plot evaluation results from JSON files', formatter_class=argparse.RawDescriptionHelpFormatter, epilog="""Examples: python plot_eval_results.py /path/to/eval1 python plot_eval_results.py Experiment1:/path/to/eval1 Experiment2:/path/to/eval2 python plot_eval_results.py /path/to/eval1 --tasks vqa gqa python plot_eval_results.py Exp1:/path/to/eval1 Exp2:/path/to/eval2 --tasks mmlu""" ) parser.add_argument('eval_folders', nargs='+', help='Evaluation folder paths, optionally with custom names (folder:name)') parser.add_argument('--tasks', default=['docvqa_val_anls', 'infovqa_val_anls', 'mme_total_score', 'mmmu_val_mmmu_acc', 'mmstar_average', 'ocrbench_ocrbench_accuracy', 'scienceqa_exact_match', 'textvqa_val_exact_match', 'average'], nargs='+', #'ai2d_exact_match', help='Specific tasks to plot (filters metrics containing these task names). Use "average_rank" for ranking-based summary metric.') parser.add_argument('--output', type=str, help='Custom filename for the saved plot (without extension)') parser.add_argument('--steps', nargs='+', type=int, help='Specific steps to plot (e.g., --steps 1000 2000 5000). If not specified, plots all available steps.') args = parser.parse_args() eval_folders = [] custom_names = [] for arg in args.eval_folders: if ':' in arg: name, folder = arg.rsplit(':', 1) eval_folders.append(folder) custom_names.append(name) else: eval_folders.append(arg) custom_names.append(None) # Check if all folders exist for eval_folder in eval_folders: if not os.path.exists(eval_folder): print(f"Error: Folder {eval_folder} does not exist") sys.exit(1) return eval_folders, custom_names, args.tasks, args.output, args.steps def main(): eval_folders, custom_names, tasks_to_plot, output_filename, steps_to_plot = parse_args() print("---------------------------") print(f"Plotting {output_filename}") # Load results from all folders all_results = [] for eval_folder in eval_folders: print(f"Loading evaluation results from: {eval_folder}") results = load_eval_results(eval_folder, tasks_to_plot) if results: print(f"Found {len(results)} evaluation steps") # Check for missing evaluation steps if specific steps are requested if steps_to_plot: available_steps = {result['step'] for result in results} missing_steps = [step for step in steps_to_plot if step not in available_steps] if missing_steps: folder_name = custom_names[eval_folders.index(eval_folder)] if custom_names and custom_names[eval_folders.index(eval_folder)] else os.path.basename(eval_folder) print(f"Warning: {folder_name} missing evaluation steps: {missing_steps}") # Check for missing evaluations if specific tasks are requested if tasks_to_plot: available_metrics = set() for result in results: available_metrics.update(k for k in result.keys() if k != 'step') missing_tasks = [] for task in tasks_to_plot: if task not in available_metrics and task not in ['average', 'average_rank']: missing_tasks.append(task) if missing_tasks: folder_name = custom_names[eval_folders.index(eval_folder)] if custom_names and custom_names[eval_folders.index(eval_folder)] else os.path.basename(eval_folder) print(f"Warning: {folder_name} does not have evaluation for tasks: {missing_tasks}") all_results.append(results) else: print(f"No evaluation results found in {eval_folder}") all_results.append([]) if any(all_results): # Compute ranking summary if requested if tasks_to_plot and 'average_rank' in tasks_to_plot: all_results = compute_ranking_summary(all_results, tasks_to_plot) plot_results(all_results, eval_folders, custom_names, tasks_to_plot, output_filename, steps_to_plot) else: print("No evaluation results found in any folder") if __name__ == "__main__": main()