"""Top-level orchestration: DatasetGenerator drives the full per-sample pipeline. For every sample this wires together, in order: 1. :class:`~src.templates.template_loader.TemplateLoader` -- pick a template + signature box. 2. :class:`~src.signatures.signature_dataset.SignatureDataset` -- pick a signature image. 3. :class:`~src.rendering.signature_preprocessor.SignaturePreprocessor` -- clean it into an RGBA stamp. 4. :class:`~src.placement.placement_engine.PlacementEngine` -- fit/rotate/position it. 5. :class:`~src.augmentation.augmentation_engine.AugmentationEngine` -- signature-layer effects. 6. :class:`~src.rendering.renderer.Renderer` -- blend onto the document. 7. :class:`~src.augmentation.augmentation_engine.AugmentationEngine` -- document-level effects. 8. :class:`~src.analysis.geometry_analyzer.GeometryAnalyzer` -- detect + measure. 9. :class:`~src.analysis.difficulty_analyzer.DifficultyAnalyzer` -- score. 10. :class:`~src.metadata.metadata_exporter.MetadataExporter` -- assemble + write JSON. 11. :class:`~src.analysis.validation_engine.ValidationEngine` -- verify. A single :class:`GeneratorContext` holds all stateless-between-calls engine instances for one process; :class:`DatasetGenerator` either drives it directly (``workers=1``) or fans it out across a :class:`concurrent.futures.ProcessPoolExecutor` pool, where each worker process builds exactly one context via ``_init_worker`` and reuses it for every sample it is assigned. Because every random draw is derived purely from ``(master_seed, sample_index, stage)``, results are identical regardless of worker count or scheduling order. """ from __future__ import annotations from concurrent.futures import ProcessPoolExecutor, as_completed from dataclasses import dataclass from pathlib import Path from typing import Any import numpy as np from tqdm import tqdm from src.analysis.difficulty_analyzer import DifficultyAnalyzer from src.analysis.geometry_analyzer import GeometryAnalyzer from src.analysis.validation_engine import ValidationEngine, ValidationReport, summarize from src.augmentation.augmentation_engine import AugmentationEngine from src.config.schema import PipelineConfig from src.metadata.metadata_exporter import MetadataExporter from src.placement.placement_engine import PlacementEngine, PlacementResult from src.rendering.printed_text import render_printed_name_stamp from src.rendering.renderer import Renderer from src.rendering.signature_preprocessor import SignaturePreprocessor from src.signatures.signature_dataset import SignatureDataset, SignatureRecord from src.templates.annotation_schema import SignatureBoxAnnotation from src.templates.template_loader import Template, TemplateLoader from src.utils.image_io import load_image_any, load_image_rgb, save_image_rgb from src.utils.logging_utils import get_logger from src.utils.names import generate_name, generate_role from src.utils.random_utils import SeedManager, resolve_scalar_or_range from src.verification.case_types import CASE_TABLE, CaseSpec from src.verification.registry import SignerIdentityRegistry from src.visualization.visualizer import save_visualization logger = get_logger(__name__) INK_COLOR_POOL: tuple[tuple[int, int, int], ...] = ( (15, 15, 20), # black (18, 22, 58), # blue-black (24, 34, 96), # ballpoint blue (12, 12, 14), # deep black ) @dataclass class SampleOutcome: """Lightweight, picklable result of generating one sample (used across process boundaries).""" sample_id: str success: bool error: "str | None" = None difficulty_tier: "str | None" = None difficulty_score: "float | None" = None validation_passed: "bool | None" = None validation_errors: int = 0 validation_warnings: int = 0 class GeneratorContext: """Holds every stateless-between-calls engine instance needed to generate a sample.""" def __init__(self, config: PipelineConfig) -> None: self.config = config self.template_loader = TemplateLoader(config.paths.image_dir, config.paths.annotations_dir) self.signature_dataset = SignatureDataset(config.paths.dataset_dir) self.preprocessor = SignaturePreprocessor() self.placement_engine = PlacementEngine(config.placement) self.augmentation_engine = AugmentationEngine(config.augmentation) self.renderer = Renderer(config.rendering) self.geometry_analyzer = GeometryAnalyzer() self.difficulty_analyzer = DifficultyAnalyzer(config.difficulty) self.metadata_exporter = MetadataExporter() self.validation_engine = ValidationEngine(config.validation, config.difficulty) # Trigger discovery eagerly so the first sample doesn't pay (and doesn't # race another worker on) the filesystem-scan + annotation-cache cost. self.template_loader.discover() self.signature_dataset.discover() self.identity_registry = SignerIdentityRegistry(self.template_loader, self.signature_dataset, config.seed) def run_sample(self, sample_index: int) -> tuple[dict[str, Any], "np.ndarray", "ValidationReport | None"]: cfg = self.config sm = SeedManager(cfg.seed, sample_index) prefix = cfg.generation.filename_prefix digits = cfg.generation.filename_digits sample_id = f"{prefix}_{sample_index:0{digits}d}" case_spec = CASE_TABLE.get(cfg.quality.case_type) if case_spec is not None: return self._run_case_sample(sample_id, sm, case_spec) return self._run_generic_sample(sample_id, sm) def _run_generic_sample(self, sample_id: str, sm: SeedManager) -> tuple[dict[str, Any], "np.ndarray", "ValidationReport | None"]: cfg = self.config template = self._select_template(sm) page_annotation = self._select_page(template) box = self._select_box(page_annotation, sm) document_image = load_image_rgb(template.image_path) canvas_h, canvas_w = document_image.shape[:2] target_box = box.as_tuple() signer_name = cfg.signer.name or generate_name(sm.python_rng("signer_name")) signer_role = cfg.signer.role or box.role or generate_role(sm.python_rng("signer_role")) signature_record = self._select_signature(sm) raw_signature = load_image_any(signature_record.path) ink_color = INK_COLOR_POOL[int(sm.numpy_rng("ink_color").integers(0, len(INK_COLOR_POOL)))] signature_rgba = self.preprocessor.process(raw_signature, ink_color=ink_color) placement = self.placement_engine.place( signature_rgba, target_box=target_box, canvas_size=(canvas_w, canvas_h), seed_manager=sm, ) signature_layer, signature_aug_records = self.augmentation_engine.apply_to_signature(placement.layer, sm) render_result = self.renderer.blend(document_image, signature_layer, sm) final_image, document_aug_records = self.augmentation_engine.apply_to_document(render_result.image, sm) geometry = self.geometry_analyzer.analyze( final_image, target_box=target_box, canvas_size=(canvas_w, canvas_h), placement_hint_bbox=placement.intended_bbox, ) difficulty = self.difficulty_analyzer.analyze(geometry) quality_rng = sm.numpy_rng("quality") quality_values = { "scanner_quality": round(resolve_scalar_or_range(cfg.quality.scanner_quality, quality_rng), 4), "paper_quality": round(resolve_scalar_or_range(cfg.quality.paper_quality, quality_rng), 4), "render_quality": round(resolve_scalar_or_range(cfg.quality.render_quality, quality_rng), 4), } metadata = self.metadata_exporter.build( document_id=sample_id, template_id=template.template_id, document_type=template.annotation.document_type, page=page_annotation.page, page_size=(canvas_w, canvas_h), dpi=page_annotation.dpi, signer_name=signer_name, signer_role=signer_role, signature_source=signature_record, placement=placement, render=render_result, signature_aug=signature_aug_records, document_aug=document_aug_records, geometry=geometry, difficulty=difficulty, quality_cfg=cfg.quality, quality_values=quality_values, seed=cfg.seed, sample_index=sm.sample_index, verification=None, ) validation_report = None if cfg.generation.run_validation: validation_report = self.validation_engine.validate(metadata) return metadata, final_image, validation_report # ------------------------------------------------------------------ # # Verification test-case generation (case_type-driven) # ------------------------------------------------------------------ # def _run_case_sample( self, sample_id: str, sm: SeedManager, case_spec: CaseSpec, ) -> tuple[dict[str, Any], "np.ndarray", "ValidationReport | None"]: cfg = self.config template = self._select_template_for_case(sm, case_spec) page_annotation = self._select_page(template) boxes = page_annotation.signature_boxes if not boxes: raise ValueError(f"Page {page_annotation.page} has no signature boxes defined.") target_box, companion_box = self._select_case_boxes(boxes, sm) document_image = load_image_rgb(template.image_path) canvas_h, canvas_w = document_image.shape[:2] target_box_tuple = target_box.as_tuple() canvas_size = (canvas_w, canvas_h) expected_group = self.identity_registry.authorized_signer_group(template.template_id, target_box.id) signer_name = cfg.signer.name or self.identity_registry.display_name(expected_group) signer_role = cfg.signer.role or target_box.role ink_color = INK_COLOR_POOL[int(sm.numpy_rng("ink_color").integers(0, len(INK_COLOR_POOL)))] document_after_companion = document_image companion_info: "dict[str, Any] | None" = None if case_spec.fill_companion: if companion_box is None: logger.warning( "Case %s wants a companion signer but template %s has only one signature box; skipping.", case_spec.code, template.template_id, ) else: companion_group = self.identity_registry.authorized_signer_group(template.template_id, companion_box.id) sm_companion = SeedManager(cfg.seed, f"{sm.sample_index}_companion") companion_record = self.signature_dataset.sample_from_group( sm_companion.numpy_rng("case:companion_signature"), companion_group, source_filter="genuine", ) companion_raw = load_image_any(companion_record.path) companion_rgba = self.preprocessor.process(companion_raw, ink_color=ink_color) companion_placement = self.placement_engine.place( companion_rgba, target_box=companion_box.as_tuple(), canvas_size=canvas_size, seed_manager=sm_companion, ) companion_layer, _ = self.augmentation_engine.apply_to_signature(companion_placement.layer, sm_companion) companion_render = self.renderer.blend(document_after_companion, companion_layer, sm_companion) document_after_companion = companion_render.image companion_info = { "box_id": companion_box.id, "filled": True, "signer_group": companion_group, "script": self.identity_registry.script_for(companion_group), "matches_expected": True, } actual_record, actual_group, actual_source_type, target_rgba = self._resolve_case_signature( case_spec, template, target_box, companion_box, expected_group, signer_name, ink_color, sm, ) if target_rgba is None: placement = PlacementResult( layer=np.zeros((canvas_h, canvas_w, 4), dtype=np.uint8), target_box=target_box_tuple, intended_bbox=target_box_tuple, x_offset=0.0, y_offset=0.0, rotation_deg=0.0, scale_factor=0.0, anchor_point=cfg.placement.anchor_point, misplaced=False, misplacement_dx=0.0, misplacement_dy=0.0, ) else: placement = self.placement_engine.place( target_rgba, target_box=target_box_tuple, canvas_size=canvas_size, seed_manager=sm, ) signature_layer, signature_aug_records = self.augmentation_engine.apply_to_signature(placement.layer, sm) render_result = self.renderer.blend(document_after_companion, signature_layer, sm) final_image, document_aug_records = self.augmentation_engine.apply_to_document(render_result.image, sm) geometry = self.geometry_analyzer.analyze( final_image, target_box=target_box_tuple, canvas_size=canvas_size, placement_hint_bbox=placement.intended_bbox, ) difficulty = self.difficulty_analyzer.analyze(geometry) quality_rng = sm.numpy_rng("quality") quality_values = { "scanner_quality": round(resolve_scalar_or_range(cfg.quality.scanner_quality, quality_rng), 4), "paper_quality": round(resolve_scalar_or_range(cfg.quality.paper_quality, quality_rng), 4), "render_quality": round(resolve_scalar_or_range(cfg.quality.render_quality, quality_rng), 4), } verification_record = { "case_type": case_spec.code, "verdict": case_spec.verdict, "reject_code": case_spec.reject_code, "target_box_id": target_box.id, "expected_signer_group": expected_group, "expected_signer_script": self.identity_registry.script_for(expected_group), "actual_signer_group": actual_group, "actual_source_type": actual_source_type, "actual_signer_script": self.identity_registry.script_for(actual_group) if actual_group else None, "companion": companion_info, } metadata = self.metadata_exporter.build( document_id=sample_id, template_id=template.template_id, document_type=template.annotation.document_type, page=page_annotation.page, page_size=(canvas_w, canvas_h), dpi=page_annotation.dpi, signer_name=signer_name, signer_role=signer_role, signature_source=actual_record, placement=placement, render=render_result, signature_aug=signature_aug_records, document_aug=document_aug_records, geometry=geometry, difficulty=difficulty, quality_cfg=cfg.quality, quality_values=quality_values, seed=cfg.seed, sample_index=sm.sample_index, verification=verification_record, ) validation_report = None if cfg.generation.run_validation: validation_report = self.validation_engine.validate(metadata) return metadata, final_image, validation_report def _select_template_for_case(self, sm: SeedManager, case_spec: CaseSpec) -> Template: cfg = self.config if cfg.document.template: return self.template_loader.get(cfg.document.template) templates = self.template_loader.discover() if case_spec.fill_companion: candidates = sorted( tid for tid, t in templates.items() if len(t.annotation.pages[0].signature_boxes) >= 2 ) if not candidates: logger.warning( "Case %s needs a multi-box template but none were found; using the full template pool.", case_spec.code, ) candidates = sorted(templates.keys()) else: candidates = sorted(templates.keys()) idx = int(sm.numpy_rng("template_selection").integers(0, len(candidates))) return self.template_loader.get(candidates[idx]) def _select_case_boxes( self, boxes: "list[SignatureBoxAnnotation]", sm: SeedManager, ) -> "tuple[SignatureBoxAnnotation, SignatureBoxAnnotation | None]": idx = int(sm.numpy_rng("box_selection").integers(0, len(boxes))) target = boxes[idx] others = [b for i, b in enumerate(boxes) if i != idx] companion = others[0] if others else None return target, companion def _resolve_case_signature( self, case_spec: CaseSpec, template: Template, target_box: "SignatureBoxAnnotation", companion_box: "SignatureBoxAnnotation | None", expected_group: str, expected_name: str, ink_color: tuple[int, int, int], sm: SeedManager, ) -> "tuple[SignatureRecord | None, str | None, str, np.ndarray | None]": """Resolve what actually goes in the target box for one case type. Returns ``(signature_record_or_none, actual_signer_group_or_none, actual_source_type, rgba_layer_or_none)``. """ mode = case_spec.signature_mode if mode == "blank": return None, None, "none", None if mode == "printed_text": layer_rgba = render_printed_name_stamp(expected_name, ink_color=ink_color) return None, None, "printed", layer_rgba if mode == "genuine_correct": group, source_filter, stage = expected_group, "genuine", "case:target_signature" elif mode == "genuine_variation": group, source_filter, stage = expected_group, "genuine", "case:target_signature_variant" elif mode == "forged_skilled": group, source_filter, stage = expected_group, "forged", "case:forged_signature" elif mode == "wrong_signer_genuine": wrong_group = self.identity_registry.random_wrong_signer_group( sm.python_rng("case:wrong_signer"), expected_group, ) group, source_filter, stage = wrong_group, "genuine", "case:wrong_signature" elif mode == "role_swap": if companion_box is None: raise ValueError(f"Case {case_spec.code} requires a companion box but none is available.") group = self.identity_registry.authorized_signer_group(template.template_id, companion_box.id) source_filter, stage = "genuine", "case:role_swap_signature" elif mode == "foreign_template_signer": group, _foreign_template_id, _foreign_box_id = self.identity_registry.foreign_authorized(template.template_id) source_filter, stage = "genuine", "case:foreign_signature" else: raise ValueError(f"Unknown signature_mode: {mode!r}") record = self.signature_dataset.sample_from_group(sm.numpy_rng(stage), group, source_filter) raw = load_image_any(record.path) layer_rgba = self.preprocessor.process(raw, ink_color=ink_color) return record, group, record.source_type, layer_rgba def _select_template(self, sm: SeedManager): cfg = self.config if cfg.document.template: return self.template_loader.get(cfg.document.template) return self.template_loader.random_template(sm.numpy_rng("template_selection")) def _select_page(self, template): cfg = self.config if cfg.document.page is not None: return template.annotation.page_by_number(cfg.document.page) return template.annotation.pages[0] def _select_box(self, page_annotation, sm: SeedManager): boxes = page_annotation.signature_boxes if not boxes: raise ValueError(f"Page {page_annotation.page} has no signature boxes defined.") idx = int(sm.numpy_rng("box_selection").integers(0, len(boxes))) return boxes[idx] def _select_signature(self, sm: SeedManager): cfg = self.config if cfg.signature_selection.signature_id: return self.signature_dataset.get(cfg.signature_selection.signature_id) return self.signature_dataset.sample( sm.numpy_rng("signature_selection"), source_filter=cfg.signature_selection.source_filter, ) _WORKER_CONTEXT: "GeneratorContext | None" = None _WORKER_CONFIG: "PipelineConfig | None" = None def _init_worker(config: PipelineConfig) -> None: global _WORKER_CONTEXT, _WORKER_CONFIG _WORKER_CONFIG = config _WORKER_CONTEXT = GeneratorContext(config) def _generate_and_export(sample_index: int) -> SampleOutcome: """Module-level worker entry point: generate one sample and write it to disk.""" assert _WORKER_CONTEXT is not None and _WORKER_CONFIG is not None, "Worker not initialized" cfg = _WORKER_CONFIG ctx = _WORKER_CONTEXT prefix = cfg.generation.filename_prefix digits = cfg.generation.filename_digits sample_id = f"{prefix}_{sample_index:0{digits}d}" output_dir = Path(cfg.paths.output_dir) png_path = output_dir / f"{sample_id}.png" json_path = output_dir / f"{sample_id}.json" if not cfg.generation.overwrite and png_path.exists() and json_path.exists(): return SampleOutcome(sample_id=sample_id, success=True, error="skipped_existing") try: metadata, final_image, validation_report = ctx.run_sample(sample_index) save_image_rgb(png_path, final_image) ctx.metadata_exporter.export(metadata, json_path) if cfg.generation.save_visualization: viz_path = output_dir / "visualizations" / f"{sample_id}_viz.png" save_visualization(viz_path, final_image, metadata) outcome = SampleOutcome( sample_id=sample_id, success=True, difficulty_tier=metadata["difficulty"]["tier"], difficulty_score=metadata["difficulty"]["score"], ) if validation_report is not None: outcome.validation_passed = validation_report.passed outcome.validation_errors = len(validation_report.errors) outcome.validation_warnings = len(validation_report.warnings) if not validation_report.passed: logger.warning( "Sample %s failed validation: %s", sample_id, "; ".join(i.message for i in validation_report.errors), ) return outcome except Exception as exc: # noqa: BLE001 - report and continue the batch logger.exception("Failed to generate sample %s", sample_id) return SampleOutcome(sample_id=sample_id, success=False, error=str(exc)) @dataclass class GenerationSummary: """Batch-level summary returned by :meth:`DatasetGenerator.generate`.""" total: int succeeded: int failed: int skipped_existing: int validation_passed: int validation_failed: int tier_counts: dict[str, int] failed_samples: list[str] def as_dict(self) -> dict[str, Any]: return { "total": self.total, "succeeded": self.succeeded, "failed": self.failed, "skipped_existing": self.skipped_existing, "validation_passed": self.validation_passed, "validation_failed": self.validation_failed, "tier_counts": self.tier_counts, "failed_samples": self.failed_samples, } class DatasetGenerator: """Drives batch generation of a synthetic signed-document dataset. Example: >>> generator = DatasetGenerator(config) >>> summary = generator.generate() """ def __init__(self, config: PipelineConfig) -> None: self.config = config Path(config.paths.output_dir).mkdir(parents=True, exist_ok=True) def generate(self) -> GenerationSummary: """Generate ``config.generation.count`` samples, sequentially or via a worker pool.""" cfg = self.config indices = list(range(cfg.generation.start_index, cfg.generation.start_index + cfg.generation.count)) outcomes: list[SampleOutcome] = [] if cfg.generation.workers <= 1: _init_worker(cfg) for idx in tqdm(indices, desc="Generating samples", unit="doc"): outcomes.append(_generate_and_export(idx)) else: with ProcessPoolExecutor( max_workers=cfg.generation.workers, initializer=_init_worker, initargs=(cfg,), ) as pool: futures = {pool.submit(_generate_and_export, idx): idx for idx in indices} for future in tqdm(as_completed(futures), total=len(futures), desc="Generating samples", unit="doc"): outcomes.append(future.result()) return self._summarize(outcomes) @staticmethod def _summarize(outcomes: list[SampleOutcome]) -> GenerationSummary: succeeded = [o for o in outcomes if o.success] failed = [o for o in outcomes if not o.success] skipped = [o for o in succeeded if o.error == "skipped_existing"] validated = [o for o in succeeded if o.validation_passed is not None] validation_passed = sum(1 for o in validated if o.validation_passed) validation_failed = sum(1 for o in validated if not o.validation_passed) tier_counts: dict[str, int] = {} for o in succeeded: if o.difficulty_tier: tier_counts[o.difficulty_tier] = tier_counts.get(o.difficulty_tier, 0) + 1 summary = GenerationSummary( total=len(outcomes), succeeded=len(succeeded), failed=len(failed), skipped_existing=len(skipped), validation_passed=validation_passed, validation_failed=validation_failed, tier_counts=tier_counts, failed_samples=[o.sample_id for o in failed], ) logger.info("Generation summary: %s", summary.as_dict()) return summary