scripts/generate_ground_truth.py

#!/usr/bin/env python3
"""Generate ground truth files (ground_truth.json, .npy) for all tasks.

Usage:
    python scripts/generate_ground_truth.py
    python scripts/generate_ground_truth.py --task task_001_flat_plate
"""
import argparse
import json
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

TASKS_ROOT = Path(__file__).parent.parent / "server" / "tasks"


def generate_one(task_dir: Path) -> dict:
    task_id = task_dir.name
    t0 = time.time()

    result = {
        "task_id": task_id,
        "success": False,
        "error": None,
    }

    try:
        from server.preprocessor import preprocess_from_code

        ref_code = (task_dir / "reference_code.py").read_text()
        gt = preprocess_from_code(ref_code, str(task_dir), task_id=task_id)

        result["success"] = True
        result["volume"] = gt.get("volume_mm3")
        result["bbox"] = gt.get("bbox_mm")
        result["euler"] = gt.get("euler_characteristic")
        result["dominant_face_type"] = gt.get("dominant_face_type")
        result["face_count"] = gt.get("face_count")

    except Exception as e:
        import traceback
        result["error"] = f"{type(e).__name__}: {e}"
        result["traceback"] = traceback.format_exc()

    result["elapsed_s"] = round(time.time() - t0, 2)
    return result


def main():
    parser = argparse.ArgumentParser(description="Generate ground truth for CadForge tasks")
    parser.add_argument("--task", type=str, default=None, help="Specific task id")
    args = parser.parse_args()

    if args.task:
        task_dirs = [TASKS_ROOT / args.task]
    else:
        task_dirs = sorted(TASKS_ROOT.glob("task_*"))

    print(f"Generating ground truth for {len(task_dirs)} tasks...")
    print("-" * 80)

    results = []
    for task_dir in task_dirs:
        r = generate_one(task_dir)
        status = "OK" if r["success"] else "FAIL"
        print(
            f"  {r['task_id']:35s} {status:5s} "
            f"vol={r.get('volume', 'N/A')} "
            f"euler={r.get('euler', 'N/A')} "
            f"faces={r.get('face_count', 'N/A')} "
            f"({r['elapsed_s']:.1f}s)"
        )
        if r["error"]:
            print(f"    ERROR: {r['error']}")
        results.append(r)

    passed = sum(1 for r in results if r["success"])
    print("-" * 80)
    print(f"Result: {passed}/{len(results)} generated successfully")

    out_path = TASKS_ROOT.parent / "verification_ground_truth.json"
    with open(out_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Saved to {out_path}")

    return 0 if passed == len(results) else 1


if __name__ == "__main__":
    sys.exit(main())