Spaces:
Runtime error
Runtime error
| #!/usr/bin/env python3 | |
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| """ | |
| Step (5) of evaluation pipeline. | |
| Validate testcase result JSONs and aggregate benchmark rows. | |
| Expected testsuite layout (aligned with evaluate_folder output): | |
| <root>/ | |
| ├── <split>/ # e.g. content, repetition | |
| │ ├── text2motion/ # text-following eval | |
| │ │ ├── overview/ # or timeline_single, timeline_multi | |
| │ │ │ └── <testcase>.json | |
| │ │ └── ... | |
| │ └── <category>/ # constraints_withtext, constraints_notext | |
| │ └── .../ # optional subdirs, e.g. root, fullbody | |
| │ └── <testcase>/ | |
| │ └── <testcase>.json | |
| Samples are discovered via rglob('meta.json') with motion.npz and gt_motion.npz in the same dir. | |
| Testcase dir = parent of a sample dir. Result file = testcase_dir.parent / f"{testcase_dir.name}.json". | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from typing import Any | |
| SPLITS = ("content", "repetition") | |
| TEXT_FOLLOWING_CATEGORIES = ("overview", "timeline_single", "timeline_multi") | |
| CONSTRAINTS_CATEGORIES = ("constraints_withtext", "constraints_notext") | |
| ROW_CATEGORIES = TEXT_FOLLOWING_CATEGORIES + CONSTRAINTS_CATEGORIES | |
| def _discover_sample_dirs(root: Path) -> list[Path]: | |
| sample_dirs: list[Path] = [] | |
| for meta_path in root.rglob("meta.json"): | |
| sample_dir = meta_path.parent | |
| if (sample_dir / "motion.npz").is_file() and (sample_dir / "gt_motion.npz").is_file(): | |
| sample_dirs.append(sample_dir) | |
| return sorted(set(sample_dirs)) | |
| def _discover_testcase_dirs(root: Path) -> list[Path]: | |
| sample_dirs = _discover_sample_dirs(root) | |
| return sorted({sample_dir.parent for sample_dir in sample_dirs}) | |
| def _expected_result_path(testcase_dir: Path) -> Path: | |
| return testcase_dir.parent / f"{testcase_dir.name}.json" | |
| def _parse_testcase_key(root: Path, testcase_dir: Path) -> tuple[str, str]: | |
| rel_parts = testcase_dir.relative_to(root).parts | |
| if len(rel_parts) < 2: | |
| raise ValueError(f"Unexpected testcase path shape: {testcase_dir} (relative: {'/'.join(rel_parts)})") | |
| split = rel_parts[0] | |
| if split not in SPLITS: | |
| raise ValueError(f"Unknown split '{split}' for testcase {testcase_dir}") | |
| if len(rel_parts) >= 3 and rel_parts[1] == "text2motion": | |
| category = rel_parts[2] | |
| if category not in TEXT_FOLLOWING_CATEGORIES: | |
| raise ValueError(f"Unknown text-following category '{category}' for testcase {testcase_dir}") | |
| else: | |
| category = rel_parts[1] | |
| if category not in CONSTRAINTS_CATEGORIES: | |
| raise ValueError(f"Unknown category '{category}' for testcase {testcase_dir}") | |
| return split, category | |
| def _accumulate_weighted( | |
| sum_acc: dict[str, float], | |
| weight_acc: dict[str, float], | |
| metric_dict: dict[str, Any], | |
| weight: float, | |
| ) -> None: | |
| for metric_name, value in metric_dict.items(): | |
| if isinstance(value, (int, float)): | |
| sum_acc[metric_name] = sum_acc.get(metric_name, 0.0) + float(value) * weight | |
| weight_acc[metric_name] = weight_acc.get(metric_name, 0.0) + weight | |
| def _to_averages( | |
| weighted_sum: dict[str, float], weight: dict[str, float] | |
| ) -> dict[str, float]: | |
| return { | |
| k: v / weight[k] | |
| for k, v in sorted(weighted_sum.items()) | |
| if weight.get(k, 0.0) > 0 | |
| } | |
| def _load_result_row( | |
| result_path: Path, | |
| ) -> tuple[float, dict[str, Any], dict[str, Any], dict[str, Any]]: | |
| payload = json.loads(result_path.read_text(encoding="utf-8")) | |
| num_motions = float(payload.get("num_motions", 1)) | |
| per_motion_mean_gen = payload.get("per_motion_mean_gen") or payload.get("per_motion_mean", {}) | |
| per_motion_mean_gt = payload.get("per_motion_mean_gt") or {} | |
| tmr = payload.get("tmr") or {} | |
| if not isinstance(per_motion_mean_gen, dict): | |
| raise ValueError(f"'per_motion_mean_gen' / 'per_motion_mean' is not a dict in {result_path}") | |
| if not isinstance(per_motion_mean_gt, dict): | |
| raise ValueError(f"'per_motion_mean_gt' is not a dict in {result_path}") | |
| if not isinstance(tmr, dict): | |
| raise ValueError(f"'tmr' is not a dict in {result_path}") | |
| return num_motions, per_motion_mean_gen, per_motion_mean_gt, tmr | |
| # Display labels for table rows (paper-style). | |
| TEXT_FOLLOWING_ROW_LABELS = { | |
| "overview": "Overview", | |
| "timeline_single": "Timeline single", | |
| "timeline_multi": "Timeline multi", | |
| } | |
| CONSTRAINTS_ROW_LABELS = { | |
| "constraints_withtext": "Constraints with text", | |
| "constraints_notext": "Constraints without text", | |
| } | |
| # Meters to cm for constraint position metrics. | |
| M_TO_CM = 100.0 | |
| def _table_value(val: float | None) -> float | str | None: | |
| """Return value for JSON table; use None for missing (omit or serialize as null).""" | |
| if val is None: | |
| return None | |
| if isinstance(val, (int, float)) and (val != val or val == float("inf")): # nan or inf | |
| return None | |
| return val | |
| def _build_tables( | |
| row_acc: dict[tuple[str, str], dict[str, Any]], | |
| ) -> dict[str, dict[str, list[dict[str, Any]]]]: | |
| """Build text_following and constraints tables per split for paper-style output.""" | |
| tables: dict[str, dict[str, list[dict[str, Any]]]] = {} | |
| for split in SPLITS: | |
| tables[split] = {"text_following": [], "constraints": []} | |
| # Text-following table: Overview, Timeline single, Timeline multi. | |
| for category in TEXT_FOLLOWING_CATEGORIES: | |
| acc = row_acc[(split, category)] | |
| per_motion_gen = _to_averages(acc["per_motion_mean_weighted_sum"], acc["per_motion_mean_weight"]) | |
| per_motion_gt = _to_averages(acc["per_motion_mean_gt_weighted_sum"], acc["per_motion_mean_gt_weight"]) | |
| tmr_avg = _to_averages(acc["tmr_weighted_sum"], acc["tmr_weight"]) if acc["tmr_weight"] else {} | |
| r03_gen = tmr_avg.get("TMR/t2m_R/R03") | |
| r03_gt = tmr_avg.get("TMR/t2m_gt_R/R03") | |
| fid_gen_text = tmr_avg.get("TMR/FID/gen_text") | |
| fid_gt_text = tmr_avg.get("TMR/FID/gt_text") | |
| fid_gen_gt = tmr_avg.get("TMR/FID/gen_gt") | |
| # Skate is velocity in m/s; convert to cm/s for display. | |
| skate_gen = per_motion_gen.get("foot_skate_from_pred_contacts") | |
| skate_gt = per_motion_gt.get("foot_skate_from_pred_contacts") | |
| contact_gen = per_motion_gen.get("foot_contact_consistency") | |
| contact_gt = per_motion_gt.get("foot_contact_consistency") | |
| row_label = TEXT_FOLLOWING_ROW_LABELS[category] | |
| tables[split]["text_following"].append( | |
| { | |
| "row": row_label, | |
| "R@3 (gen)": _table_value(r03_gen), | |
| "R@3 (GT)": _table_value(r03_gt), | |
| "FID gen-text": _table_value(fid_gen_text), | |
| "FID GT-text": _table_value(fid_gt_text), | |
| "FID gen-GT": _table_value(fid_gen_gt), | |
| "Skate (gen, cm/s)": _table_value(skate_gen * 100.0 if skate_gen is not None else None), | |
| "Skate (GT, cm/s)": _table_value(skate_gt * 100.0 if skate_gt is not None else None), | |
| "Contact (gen)": _table_value(contact_gen), | |
| "Contact (GT)": _table_value(contact_gt), | |
| } | |
| ) | |
| # Constraints table: Constraints with text, Constraints without text. | |
| for category in CONSTRAINTS_CATEGORIES: | |
| acc = row_acc[(split, category)] | |
| per_motion_gen = _to_averages(acc["per_motion_mean_weighted_sum"], acc["per_motion_mean_weight"]) | |
| per_motion_gt = _to_averages(acc["per_motion_mean_gt_weighted_sum"], acc["per_motion_mean_gt_weight"]) | |
| row_label = CONSTRAINTS_ROW_LABELS[category] | |
| row_dict: dict[str, Any] = { | |
| "row": row_label, | |
| "Full-Body Pos (gen, cm)": _table_value( | |
| per_motion_gen.get("constraint_fullbody_keyframe") * M_TO_CM | |
| if per_motion_gen.get("constraint_fullbody_keyframe") is not None | |
| else None | |
| ), | |
| "Full-Body Pos (GT, cm)": _table_value( | |
| per_motion_gt.get("constraint_fullbody_keyframe") * M_TO_CM | |
| if per_motion_gt.get("constraint_fullbody_keyframe") is not None | |
| else None | |
| ), | |
| "End-Effector Pos (gen, cm)": _table_value( | |
| per_motion_gen.get("constraint_end_effector") * M_TO_CM | |
| if per_motion_gen.get("constraint_end_effector") is not None | |
| else None | |
| ), | |
| "End-Effector Pos (GT, cm)": _table_value( | |
| per_motion_gt.get("constraint_end_effector") * M_TO_CM | |
| if per_motion_gt.get("constraint_end_effector") is not None | |
| else None | |
| ), | |
| "End-Effector Rot (deg)": None, # Not implemented in metrics. | |
| "2D Root Pos (gen, cm)": _table_value( | |
| per_motion_gen.get("constraint_root2d_err") * M_TO_CM | |
| if per_motion_gen.get("constraint_root2d_err") is not None | |
| else None | |
| ), | |
| "2D Root Pos (GT, cm)": _table_value( | |
| per_motion_gt.get("constraint_root2d_err") * M_TO_CM | |
| if per_motion_gt.get("constraint_root2d_err") is not None | |
| else None | |
| ), | |
| "2D Pelvis Pos@95% (gen, cm)": _table_value( | |
| per_motion_gen.get("constraint_root2d_err_p95") * M_TO_CM | |
| if per_motion_gen.get("constraint_root2d_err_p95") is not None | |
| else None | |
| ), | |
| "2D Pelvis Pos@95% (GT, cm)": _table_value( | |
| per_motion_gt.get("constraint_root2d_err_p95") * M_TO_CM | |
| if per_motion_gt.get("constraint_root2d_err_p95") is not None | |
| else None | |
| ), | |
| } | |
| tables[split]["constraints"].append(row_dict) | |
| return tables | |
| def _fmt_md(val: float | None, decimals: int) -> str: | |
| """Format a numeric value for a markdown cell, or '-' for None/NaN.""" | |
| if val is None: | |
| return "-" | |
| if isinstance(val, float) and (val != val or val == float("inf")): | |
| return "-" | |
| return f"{val:.{decimals}f}" | |
| def _print_tf_formatted_md( | |
| splits_data: list[tuple[str, list[dict[str, Any]]]], | |
| title: str, | |
| ) -> None: | |
| """Print text-following table in markdown, mirroring the terminal layout.""" | |
| groups = ["Overview", "Timeline single", "Timeline multi"] | |
| specs: list[tuple[str, int]] = [ | |
| ("R@3\u2191", 2), | |
| ("FID\u2193", 3), | |
| ("Skate\u2193", 3), | |
| ("Contact\u2191", 3), | |
| ] | |
| gt_keys = ["R@3 (GT)", None, "Skate (GT, cm/s)", "Contact (GT)"] | |
| gen_keys = ["R@3 (gen)", "FID gen-GT", "Skate (gen, cm/s)", "Contact (gen)"] | |
| gt_defaults: list[float | None] = [None, 0.0, None, None] | |
| headers = [""] | |
| for g in groups: | |
| for hdr, _ in specs: | |
| headers.append(f"{g} {hdr}") | |
| print(f"\n### {title}\n") | |
| print("| " + " | ".join(headers) + " |") | |
| print("| " + " | ".join("---" for _ in headers) + " |") | |
| for split_label, rows in splits_data: | |
| for row_type, keys, defaults in [ | |
| ("Ground Truth", gt_keys, gt_defaults), | |
| ("Method", gen_keys, [None] * len(specs)), | |
| ]: | |
| cells = [f"**{split_label}** {row_type}"] | |
| for row in rows: | |
| for j, (_, dec) in enumerate(specs): | |
| key = keys[j] | |
| val = defaults[j] if key is None else row.get(key) | |
| cells.append(_fmt_md(val, dec)) | |
| print("| " + " | ".join(cells) + " |") | |
| print() | |
| def _print_c_formatted_md( | |
| splits_data: list[tuple[str, list[dict[str, Any]]]], | |
| title: str, | |
| ) -> None: | |
| """Print constraints table in markdown, mirroring the terminal layout.""" | |
| groups = ["With text", "Without text"] | |
| specs: list[tuple[str, int]] = [ | |
| ("FB Pos\u2193", 3), | |
| ("EE Pos\u2193", 3), | |
| ("EE Rot\u2193", 3), | |
| ("2D Root\u2193", 3), | |
| ("Pelvis@95%", 2), | |
| ] | |
| gt_keys = [ | |
| "Full-Body Pos (GT, cm)", | |
| "End-Effector Pos (GT, cm)", | |
| "End-Effector Rot (deg)", | |
| "2D Root Pos (GT, cm)", | |
| "2D Pelvis Pos@95% (GT, cm)", | |
| ] | |
| gen_keys = [ | |
| "Full-Body Pos (gen, cm)", | |
| "End-Effector Pos (gen, cm)", | |
| "End-Effector Rot (deg)", | |
| "2D Root Pos (gen, cm)", | |
| "2D Pelvis Pos@95% (gen, cm)", | |
| ] | |
| headers = [""] | |
| for g in groups: | |
| for hdr, _ in specs: | |
| headers.append(f"{g} {hdr}") | |
| print(f"\n### {title}\n") | |
| print("| " + " | ".join(headers) + " |") | |
| print("| " + " | ".join("---" for _ in headers) + " |") | |
| for split_label, rows in splits_data: | |
| for row_type, keys in [("Ground Truth", gt_keys), ("Method", gen_keys)]: | |
| cells = [f"**{split_label}** {row_type}"] | |
| for row in rows: | |
| for j, (_, dec) in enumerate(specs): | |
| cells.append(_fmt_md(row.get(keys[j]), dec)) | |
| print("| " + " | ".join(cells) + " |") | |
| print() | |
| def _print_formatted_gt_method_md( | |
| tables: dict[str, dict[str, list[dict[str, Any]]]], | |
| ) -> None: | |
| """Print combined tables in markdown format, mirroring the terminal layout.""" | |
| tf_splits: list[tuple[str, list[dict[str, Any]]]] = [] | |
| c_splits: list[tuple[str, list[dict[str, Any]]]] = [] | |
| for split in SPLITS: | |
| split_tables = tables.get(split, {}) | |
| tf_rows = split_tables.get("text_following", []) | |
| c_rows = split_tables.get("constraints", []) | |
| if tf_rows and len(tf_rows) == 3: | |
| tf_splits.append((split.capitalize(), tf_rows)) | |
| if c_rows and len(c_rows) == 2: | |
| c_splits.append((split.capitalize(), c_rows)) | |
| if tf_splits: | |
| _print_tf_formatted_md(tf_splits, "Text-Following Evaluation") | |
| if c_splits: | |
| _print_c_formatted_md(c_splits, "Constrained Evaluation") | |
| def _fmt(val: float | None, decimals: int, width: int) -> str: | |
| """Format a numeric value right-aligned to *width*, or '-' for None.""" | |
| if val is None: | |
| return f"{'-':>{width}}" | |
| return f"{val:>{width}.{decimals}f}" | |
| def _print_grouped_rows( | |
| label: str, | |
| rows: list[dict[str, Any]], | |
| specs: list[tuple[str, int, int]], | |
| keys: list[str], | |
| mw: int, | |
| sep: str, | |
| ) -> None: | |
| """Print one data row across all column groups.""" | |
| parts = [f"{label:<{mw}}"] | |
| for i, row in enumerate(rows): | |
| if i: | |
| parts.append(sep) | |
| for j, (_, dec, w) in enumerate(specs): | |
| parts.append(_fmt(row.get(keys[j]), dec, w)) | |
| print("".join(parts)) | |
| def _print_tf_formatted( | |
| splits_data: list[tuple[str, list[dict[str, Any]]]], | |
| title: str, | |
| ) -> None: | |
| """Print text-following table with Overview / Timeline single / Timeline multi groups. | |
| *splits_data* is a list of ``(split_label, category_rows)`` tuples so | |
| that content and repetition splits appear as separate row-pairs inside | |
| one table. | |
| """ | |
| groups = ["Overview", "Timeline single", "Timeline multi"] | |
| specs: list[tuple[str, int, int]] = [ | |
| ("R@3\u2191", 2, 7), | |
| ("FID\u2193", 3, 7), | |
| ("Skate\u2193", 3, 9), | |
| ("Contact\u2191", 3, 10), | |
| ] | |
| gt_keys = ["R@3 (GT)", None, "Skate (GT, cm/s)", "Contact (GT)"] | |
| gen_keys = ["R@3 (gen)", "FID gen-GT", "Skate (gen, cm/s)", "Contact (gen)"] | |
| gt_defaults: list[float | None] = [None, 0.0, None, None] | |
| mw = 16 | |
| gw = sum(s[2] for s in specs) | |
| sep = " | " | |
| total_w = mw + len(groups) * gw + (len(groups) - 1) * len(sep) | |
| print(f"\n{title:^{total_w}}") | |
| print("=" * total_w) | |
| parts: list[str] = [" " * mw] | |
| for i, g in enumerate(groups): | |
| if i: | |
| parts.append(sep) | |
| parts.append(g.center(gw)) | |
| print("".join(parts)) | |
| parts = [f"{'':<{mw}}"] | |
| for i in range(len(groups)): | |
| if i: | |
| parts.append(sep) | |
| for hdr, _, w in specs: | |
| parts.append(f"{hdr:>{w}}") | |
| print("".join(parts)) | |
| parts = ["\u2500" * mw] | |
| for i in range(len(groups)): | |
| if i: | |
| parts.append("\u2500\u253c\u2500") | |
| parts.append("\u2500" * gw) | |
| print("".join(parts)) | |
| for si, (split_label, rows) in enumerate(splits_data): | |
| tag = f"\u2500\u2500 {split_label} " | |
| print(tag + "\u2500" * (total_w - len(tag))) | |
| parts = [f"{'Ground Truth':<{mw}}"] | |
| for i, row in enumerate(rows): | |
| if i: | |
| parts.append(sep) | |
| for j, (_, dec, w) in enumerate(specs): | |
| key = gt_keys[j] | |
| val = gt_defaults[j] if key is None else row.get(key) | |
| parts.append(_fmt(val, dec, w)) | |
| print("".join(parts)) | |
| _print_grouped_rows("Method", rows, specs, gen_keys, mw, sep) | |
| print() | |
| def _print_c_formatted( | |
| splits_data: list[tuple[str, list[dict[str, Any]]]], | |
| title: str, | |
| ) -> None: | |
| """Print constraints table with With text / Without text groups. | |
| *splits_data* is a list of ``(split_label, category_rows)`` tuples. | |
| """ | |
| groups = ["With text", "Without text"] | |
| specs: list[tuple[str, int, int]] = [ | |
| ("FB Pos\u2193", 3, 10), | |
| ("EE Pos\u2193", 3, 10), | |
| ("EE Rot\u2193", 3, 10), | |
| ("2D Root\u2193", 3, 11), | |
| ("Pelvis@95%", 2, 12), | |
| ] | |
| gt_keys = [ | |
| "Full-Body Pos (GT, cm)", | |
| "End-Effector Pos (GT, cm)", | |
| "End-Effector Rot (deg)", | |
| "2D Root Pos (GT, cm)", | |
| "2D Pelvis Pos@95% (GT, cm)", | |
| ] | |
| gen_keys = [ | |
| "Full-Body Pos (gen, cm)", | |
| "End-Effector Pos (gen, cm)", | |
| "End-Effector Rot (deg)", | |
| "2D Root Pos (gen, cm)", | |
| "2D Pelvis Pos@95% (gen, cm)", | |
| ] | |
| mw = 16 | |
| gw = sum(s[2] for s in specs) | |
| sep = " | " | |
| total_w = mw + len(groups) * gw + (len(groups) - 1) * len(sep) | |
| print(f"\n{title:^{total_w}}") | |
| print("=" * total_w) | |
| parts: list[str] = [" " * mw] | |
| for i, g in enumerate(groups): | |
| if i: | |
| parts.append(sep) | |
| parts.append(g.center(gw)) | |
| print("".join(parts)) | |
| parts = [f"{'':<{mw}}"] | |
| for i in range(len(groups)): | |
| if i: | |
| parts.append(sep) | |
| for hdr, _, w in specs: | |
| parts.append(f"{hdr:>{w}}") | |
| print("".join(parts)) | |
| parts = ["\u2500" * mw] | |
| for i in range(len(groups)): | |
| if i: | |
| parts.append("\u2500\u253c\u2500") | |
| parts.append("\u2500" * gw) | |
| print("".join(parts)) | |
| for si, (split_label, rows) in enumerate(splits_data): | |
| tag = f"\u2500\u2500 {split_label} " | |
| print(tag + "\u2500" * (total_w - len(tag))) | |
| _print_grouped_rows("Ground Truth", rows, specs, gt_keys, mw, sep) | |
| _print_grouped_rows("Method", rows, specs, gen_keys, mw, sep) | |
| print() | |
| def _print_formatted_gt_method( | |
| tables: dict[str, dict[str, list[dict[str, Any]]]], | |
| ) -> None: | |
| """Print combined tables with column groups separated by vertical bars. | |
| Content and repetition splits are shown as separate row-pairs inside one text-following table | |
| and one constraints table. | |
| """ | |
| tf_splits: list[tuple[str, list[dict[str, Any]]]] = [] | |
| c_splits: list[tuple[str, list[dict[str, Any]]]] = [] | |
| for split in SPLITS: | |
| split_tables = tables.get(split, {}) | |
| tf_rows = split_tables.get("text_following", []) | |
| c_rows = split_tables.get("constraints", []) | |
| if tf_rows and len(tf_rows) == 3: | |
| tf_splits.append((split.capitalize(), tf_rows)) | |
| if c_rows and len(c_rows) == 2: | |
| c_splits.append((split.capitalize(), c_rows)) | |
| if tf_splits: | |
| _print_tf_formatted(tf_splits, "Text-Following Evaluation") | |
| if c_splits: | |
| _print_c_formatted(c_splits, "Constrained Evaluation") | |
| def _build_summary(root: Path) -> dict[str, Any]: | |
| testcase_dirs = _discover_testcase_dirs(root) | |
| if not testcase_dirs: | |
| raise SystemExit( | |
| f"No testcase folders found under {root} (expected folders containing meta.json + motion.npz + gt_motion.npz samples)." | |
| ) | |
| missing_results: list[Path] = [] | |
| for testcase_dir in testcase_dirs: | |
| result_path = _expected_result_path(testcase_dir) | |
| if not result_path.is_file(): | |
| missing_results.append(result_path) | |
| if missing_results: | |
| missing_text = "\n".join(str(path) for path in missing_results) | |
| raise SystemExit(f"Missing {len(missing_results)} testcase result JSON files:\n{missing_text}") | |
| row_acc: dict[tuple[str, str], dict[str, Any]] = {} | |
| for split in SPLITS: | |
| for category in ROW_CATEGORIES: | |
| row_acc[(split, category)] = { | |
| "num_testcases": 0, | |
| "num_motions": 0.0, | |
| "per_motion_mean_weighted_sum": {}, | |
| "per_motion_mean_weight": {}, | |
| "per_motion_mean_gt_weighted_sum": {}, | |
| "per_motion_mean_gt_weight": {}, | |
| "tmr_weighted_sum": {}, | |
| "tmr_weight": {}, | |
| } | |
| for testcase_dir in testcase_dirs: | |
| split, category = _parse_testcase_key(root, testcase_dir) | |
| result_path = _expected_result_path(testcase_dir) | |
| num_motions, per_motion_mean_gen, per_motion_mean_gt, tmr = _load_result_row(result_path) | |
| acc = row_acc[(split, category)] | |
| acc["num_testcases"] += 1 | |
| acc["num_motions"] += num_motions | |
| _accumulate_weighted( | |
| acc["per_motion_mean_weighted_sum"], | |
| acc["per_motion_mean_weight"], | |
| per_motion_mean_gen, | |
| num_motions, | |
| ) | |
| if per_motion_mean_gt: | |
| _accumulate_weighted( | |
| acc["per_motion_mean_gt_weighted_sum"], | |
| acc["per_motion_mean_gt_weight"], | |
| per_motion_mean_gt, | |
| num_motions, | |
| ) | |
| if tmr: | |
| _accumulate_weighted( | |
| acc["tmr_weighted_sum"], | |
| acc["tmr_weight"], | |
| tmr, | |
| num_motions, | |
| ) | |
| rows: list[dict[str, Any]] = [] | |
| for split in SPLITS: | |
| for category in ROW_CATEGORIES: | |
| acc = row_acc[(split, category)] | |
| tmr_avg = _to_averages(acc["tmr_weighted_sum"], acc["tmr_weight"]) if acc["tmr_weight"] else {} | |
| per_motion_gt_avg = _to_averages(acc["per_motion_mean_gt_weighted_sum"], acc["per_motion_mean_gt_weight"]) | |
| row_dict: dict[str, Any] = { | |
| "split": split, | |
| "category": category, | |
| "num_testcases": acc["num_testcases"], | |
| "num_motions": int(acc["num_motions"]), | |
| "per_motion_mean": _to_averages(acc["per_motion_mean_weighted_sum"], acc["per_motion_mean_weight"]), | |
| "tmr": tmr_avg, | |
| } | |
| if per_motion_gt_avg: | |
| row_dict["per_motion_mean_gt"] = per_motion_gt_avg | |
| rows.append(row_dict) | |
| # Combined constraints row for this split. | |
| withtext = row_acc[(split, "constraints_withtext")] | |
| notext = row_acc[(split, "constraints_notext")] | |
| combined_per_motion = defaultdict(float) | |
| combined_per_motion_weight = defaultdict(float) | |
| combined_per_motion_gt = defaultdict(float) | |
| combined_per_motion_gt_weight = defaultdict(float) | |
| combined_tmr = defaultdict(float) | |
| combined_tmr_weight = defaultdict(float) | |
| for sum_key, weight_key, sum_acc, weight_acc in ( | |
| ("per_motion_mean_weighted_sum", "per_motion_mean_weight", combined_per_motion, combined_per_motion_weight), | |
| ("per_motion_mean_gt_weighted_sum", "per_motion_mean_gt_weight", combined_per_motion_gt, combined_per_motion_gt_weight), | |
| ("tmr_weighted_sum", "tmr_weight", combined_tmr, combined_tmr_weight), | |
| ): | |
| for src in (withtext, notext): | |
| for k, v in src[sum_key].items(): | |
| sum_acc[k] += v | |
| for k, w in src[weight_key].items(): | |
| weight_acc[k] += w | |
| combined_tmr_avg = _to_averages(dict(combined_tmr), dict(combined_tmr_weight)) if combined_tmr_weight else {} | |
| combined_gt_avg = _to_averages(dict(combined_per_motion_gt), dict(combined_per_motion_gt_weight)) | |
| combined_row: dict[str, Any] = { | |
| "split": split, | |
| "category": "constraints", | |
| "num_testcases": withtext["num_testcases"] + notext["num_testcases"], | |
| "num_motions": int(withtext["num_motions"] + notext["num_motions"]), | |
| "per_motion_mean": _to_averages(dict(combined_per_motion), dict(combined_per_motion_weight)), | |
| "tmr": combined_tmr_avg, | |
| } | |
| if combined_gt_avg: | |
| combined_row["per_motion_mean_gt"] = combined_gt_avg | |
| rows.append(combined_row) | |
| tables = _build_tables(row_acc) | |
| return { | |
| "folder": str(root), | |
| "num_testcases": len(testcase_dirs), | |
| "rows": rows, | |
| "tables": tables, | |
| } | |
| def main() -> None: | |
| parser = argparse.ArgumentParser( | |
| description=("Validate testcase XXX.json result files and aggregate averages by split/category.") | |
| ) | |
| parser.add_argument( | |
| "folder", | |
| type=Path, | |
| help="Testsuite root folder (contains content/ and repetition/).", | |
| ) | |
| parser.add_argument( | |
| "--output", | |
| type=Path, | |
| default=None, | |
| help="Optional output JSON path. Default: <folder>/summary_rows.json", | |
| ) | |
| parser.add_argument( | |
| "--format", | |
| choices=["terminal", "md"], | |
| default="terminal", | |
| dest="table_format", | |
| help="Table output format: 'terminal' (default) for fixed-width tables, 'md' for markdown.", | |
| ) | |
| args = parser.parse_args() | |
| folder = args.folder.resolve() | |
| if not folder.is_dir(): | |
| raise SystemExit(f"Folder does not exist: {folder}") | |
| summary = _build_summary(folder) | |
| out_path = args.output.resolve() if args.output else folder / "summary_rows.json" | |
| out_path.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") | |
| print(f"Wrote aggregated summary: {out_path}") | |
| print(f"Rows: {len(summary['rows'])}, testcases: {summary['num_testcases']}") | |
| if args.table_format == "md": | |
| _print_formatted_gt_method_md(summary["tables"]) | |
| else: | |
| _print_formatted_gt_method(summary["tables"]) | |
| if __name__ == "__main__": | |
| main() | |