File size: 4,855 Bytes
97b64b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | #!/usr/bin/env python
"""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())
|