| |
| """CLI entry point for the synthetic signed bank document generation pipeline. |
| |
| Examples: |
| Generate 10 samples with all defaults:: |
| |
| python main.py |
| |
| Generate 10,000 samples using a custom config:: |
| |
| python main.py --config configs/generate.yaml --count 10000 |
| |
| Override arbitrary nested config values from the command line:: |
| |
| python main.py --set placement.rotation_deg.max=6 --set augmentation.coffee_stain.enabled=true |
| |
| Generate a single, fully reproducible sample from one template with a debug figure:: |
| |
| python main.py --count 1 --seed 42 --template bank_doc_03 --save-visualization |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import logging |
| import sys |
| from pathlib import Path |
|
|
| from src.config.loader import load_config, save_resolved_config |
| from src.generator import DatasetGenerator |
| from src.utils.logging_utils import configure_logging, get_logger |
|
|
| logger = get_logger(__name__) |
|
|
|
|
| def build_arg_parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser( |
| description="Generate synthetic signed bank documents with rich ML-ready metadata.", |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| ) |
| parser.add_argument("--config", type=str, default=None, help="Path to a user YAML config, merged over configs/default.yaml.") |
| parser.add_argument("--count", type=int, default=None, help="Number of samples to generate (overrides config).") |
| parser.add_argument("--seed", type=int, default=None, help="Master random seed (overrides config).") |
| parser.add_argument("--workers", type=int, default=None, help="Number of worker processes (overrides config).") |
| parser.add_argument("--output", type=str, default=None, help="Output directory (overrides config).") |
| parser.add_argument("--template", type=str, default=None, help="Restrict generation to a single template id, e.g. bank_doc_03.") |
| parser.add_argument("--start-index", type=int, default=None, help="First sample index (overrides config).") |
| parser.add_argument("--overwrite", action="store_true", default=None, help="Regenerate samples even if their output files already exist.") |
| parser.add_argument("--save-visualization", action="store_true", default=None, help="Also save a debug visualization PNG per sample.") |
| parser.add_argument( |
| "--set", dest="overrides", action="append", default=[], |
| metavar="key.path=value", help="Arbitrary dot-path config override; may be repeated.", |
| ) |
| parser.add_argument("--log-level", type=str, default="INFO", choices=["DEBUG", "INFO", "WARNING", "ERROR"]) |
| return parser |
|
|
|
|
| def _collect_overrides(args: argparse.Namespace) -> list[str]: |
| overrides = list(args.overrides) |
| if args.count is not None: |
| overrides.append(f"generation.count={args.count}") |
| if args.workers is not None: |
| overrides.append(f"generation.workers={args.workers}") |
| if args.output is not None: |
| overrides.append(f"paths.output_dir={args.output}") |
| if args.start_index is not None: |
| overrides.append(f"generation.start_index={args.start_index}") |
| if args.overwrite: |
| overrides.append("generation.overwrite=true") |
| if args.save_visualization: |
| overrides.append("generation.save_visualization=true") |
| if args.seed is not None: |
| overrides.append(f"seed={args.seed}") |
| if args.template is not None: |
| overrides.append(f"document.template={args.template}") |
| return overrides |
|
|
|
|
| def main(argv: "list[str] | None" = None) -> int: |
| args = build_arg_parser().parse_args(argv) |
| configure_logging(level=getattr(logging, args.log_level)) |
|
|
| overrides = _collect_overrides(args) |
| config = load_config(config_path=args.config, overrides=overrides) |
|
|
| logger.info("Starting generation: seed=%d, count=%d, workers=%d, output=%s", |
| config.seed, config.generation.count, config.generation.workers, config.paths.output_dir) |
|
|
| output_dir = Path(config.paths.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
| save_resolved_config(config, output_dir / "resolved_config.yaml") |
|
|
| generator = DatasetGenerator(config) |
| summary = generator.generate() |
|
|
| summary_path = output_dir / "generation_summary.json" |
| summary_path.write_text(json.dumps(summary.as_dict(), indent=2), encoding="utf-8") |
| logger.info("Wrote generation summary to %s", summary_path) |
| logger.info( |
| "Done: %d/%d succeeded, %d failed, %d skipped (already existed). Difficulty tiers: %s", |
| summary.succeeded, summary.total, summary.failed, summary.skipped_existing, summary.tier_counts, |
| ) |
| if summary.validation_failed: |
| logger.warning("%d sample(s) failed validation.", summary.validation_failed) |
|
|
| return 1 if summary.failed else 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
