Rich Document Image Generation
Download Recipe
1 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 2 # SPDX-License-Identifier: Apache-2.0 3 # /// script 4 # requires-python = ">=3.10" 5 # dependencies = [ 6 # "data-designer", 7 # "pandas", 8 # "pyarrow", 9 # ] 10 # /// 11 """Rich Document Image Generation Recipe 12 13 Generate synthetic business-document page images with controlled variation. 14 Each generated row pairs an image with the metadata that produced it, making 15 the output useful as seed data for visual QA, OCR robustness, multimodal 16 judging, and document-understanding experiments. 17 18 Prerequisites: 19 - An image-generation provider key for the selected model. The defaults use 20 OpenRouter, so set OPENROUTER_API_KEY before running. 21 22 Run: 23 # Generate 5 rich document images with the default OpenRouter model. 24 uv run rich_document_images.py --num-records 5 25 26 # Export a VQA-ready seed parquet with base64 image bytes and image metadata. 27 uv run rich_document_images.py --num-records 25 --export-seed rich_document_seed.parquet 28 29 # Use a different provider or image model. 30 uv run rich_document_images.py --model-provider openrouter --model-id google/gemini-3.1-flash-image-preview 31 """ 32 33 from __future__ import annotations 34 35 import argparse 36 import base64 37 from collections.abc import Iterable 38 from pathlib import Path 39 40 import pandas as pd 41 from PIL import Image 42 43 import data_designer.config as dd 44 from data_designer.interface import DataDesigner, DatasetCreationResults 45 46 DEFAULT_MODEL_PROVIDER = "openrouter" 47 DEFAULT_MODEL_ID = "google/gemini-3.1-flash-image-preview" 48 DEFAULT_MODEL_ALIAS = "document-generation-model" 49 50 SEED_METADATA_COLUMNS = [ 51 "document_type", 52 "primary_visual", 53 "secondary_visual", 54 "layout_style", 55 "document_condition", 56 ] 57 58 59 def build_model_configs( 60 *, 61 model_provider: str, 62 model_id: str, 63 model_alias: str, 64 image_size: str, 65 aspect_ratio: str, 66 max_parallel_requests: int, 67 ) -> list[dd.ModelConfig]: 68 """Build a provider-agnostic image-generation model config.""" 69 return [ 70 dd.ModelConfig( 71 alias=model_alias, 72 model=model_id, 73 provider=model_provider, 74 inference_parameters=dd.ImageInferenceParams( 75 extra_body={ 76 "modalities": ["image", "text"], 77 "image_config": { 78 "aspect_ratio": aspect_ratio, 79 "image_size": image_size, 80 }, 81 }, 82 max_parallel_requests=max_parallel_requests, 83 ), 84 skip_health_check=True, 85 ) 86 ] 87 88 89 def add_category( 90 config_builder: dd.DataDesignerConfigBuilder, 91 name: str, 92 values: list[str], 93 weights: list[float] | None = None, 94 ) -> None: 95 """Add a categorical sampler column.""" 96 config_builder.add_column( 97 dd.SamplerColumnConfig( 98 name=name, 99 sampler_type=dd.SamplerType.CATEGORY, 100 params=dd.CategorySamplerParams(values=values, weights=weights), 101 ) 102 ) 103 104 105 def add_visual_variation_id(config_builder: dd.DataDesignerConfigBuilder) -> None: 106 """Add a unique row-level key that discourages duplicate image generations.""" 107 config_builder.add_column( 108 dd.SamplerColumnConfig( 109 name="visual_variation_id", 110 sampler_type=dd.SamplerType.UUID, 111 params=dd.UUIDSamplerParams(prefix="doc-", short_form=True), 112 ) 113 ) 114 115 116 def build_config( 117 *, 118 model_provider: str = DEFAULT_MODEL_PROVIDER, 119 model_id: str = DEFAULT_MODEL_ID, 120 model_alias: str = DEFAULT_MODEL_ALIAS, 121 image_size: str = "1K", 122 aspect_ratio: str = "2:3", 123 max_parallel_requests: int = 10, 124 ) -> dd.DataDesignerConfigBuilder: 125 """Build a rich document image-generation pipeline.""" 126 model_configs = build_model_configs( 127 model_provider=model_provider, 128 model_id=model_id, 129 model_alias=model_alias, 130 image_size=image_size, 131 aspect_ratio=aspect_ratio, 132 max_parallel_requests=max_parallel_requests, 133 ) 134 config_builder = dd.DataDesignerConfigBuilder(model_configs=model_configs) 135 add_visual_variation_id(config_builder) 136 137 add_category( 138 config_builder, 139 "document_type", 140 [ 141 "quarterly business review", 142 "market research brief", 143 "operations dashboard export", 144 "clinical trial status report", 145 "sustainability impact report", 146 "financial variance memo", 147 "customer support incident review", 148 "supply chain risk assessment", 149 "product launch readiness plan", 150 "employee engagement summary", 151 ], 152 weights=[0.12, 0.10, 0.14, 0.08, 0.08, 0.12, 0.12, 0.10, 0.12, 0.12], 153 ) 154 155 add_category( 156 config_builder, 157 "organization_name", 158 [ 159 "Aster Analytics", 160 "Blue Ridge Health", 161 "CedarWorks Manufacturing", 162 "DeltaGrid Energy", 163 "Evergreen Mobility", 164 "Harborlight Retail", 165 "Northstar Robotics", 166 "Redwood BioSystems", 167 "Summit Cloud Services", 168 "Valley Forge Logistics", 169 ], 170 ) 171 172 add_category( 173 config_builder, 174 "document_owner", 175 [ 176 "Maya Chen", 177 "Jonas Patel", 178 "Elena Garcia", 179 "Noah Williams", 180 "Amara Okafor", 181 "Theo Martin", 182 "Priya Raman", 183 "Sofia Rossi", 184 "Lena Fischer", 185 "Caleb Brooks", 186 ], 187 ) 188 189 add_category( 190 config_builder, 191 "owner_role", 192 [ 193 "VP Operations", 194 "Finance Director", 195 "Clinical Program Manager", 196 "Customer Success Lead", 197 "Risk Officer", 198 "Product Launch Owner", 199 "People Analytics Partner", 200 ], 201 ) 202 203 add_category( 204 config_builder, 205 "audience", 206 [ 207 "executive leadership", 208 "finance review committee", 209 "field operations managers", 210 "clinical program leads", 211 "board audit committee", 212 "customer success directors", 213 ], 214 ) 215 216 add_category( 217 config_builder, 218 "content_theme", 219 [ 220 "quarterly revenue performance and forecast variance", 221 "regional customer adoption and churn risk", 222 "service-level agreement compliance and incident aging", 223 "inventory throughput, backorders, and supplier delays", 224 "trial enrollment, site activation, and adverse event counts", 225 "energy consumption, emissions, and sustainability targets", 226 "hiring funnel conversion, offer acceptance, and attrition", 227 "product launch milestones, owners, and readiness status", 228 ], 229 ) 230 231 add_category( 232 config_builder, 233 "primary_visual", 234 [ 235 "clustered bar chart comparing three regions across four quarters", 236 "line chart with two series, annotated inflection points, and a target band", 237 "stacked area chart showing category mix over six months", 238 "waterfall chart showing contributors to budget variance", 239 "scatter plot with labeled outliers and a trend line", 240 "Gantt-style timeline with milestones and owner initials", 241 "heatmap matrix with risk severity by team and region", 242 "donut chart with callout labels and percentages", 243 ], 244 ) 245 246 add_category( 247 config_builder, 248 "secondary_visual", 249 [ 250 "dense financial table with subtotals and variance arrows", 251 "KPI card row with current value, target, delta, and traffic-light status", 252 "two-column risk register with owner, due date, and mitigation note", 253 "small process diagram with arrows between four labeled stages", 254 "ranked list table with sparklines in the final column", 255 "compact map inset with region labels and numeric badges", 256 "executive callout box with three bullet conclusions", 257 "signature block plus approval checklist", 258 ], 259 ) 260 261 add_category( 262 config_builder, 263 "layout_style", 264 [ 265 "clean consulting report page with narrow margins and section dividers", 266 "dashboard export with a top filter bar and grid of panels", 267 "formal memo with letterhead, dense paragraphs, and one embedded chart", 268 "board-pack page with title ribbon, footnotes, and small-print source notes", 269 "compliance form with checkboxes, tables, and stamped approval", 270 "research brief with abstract, sidebar definitions, and figure captions", 271 "operations one-pager with color-coded status chips and action table", 272 ], 273 ) 274 275 add_category( 276 config_builder, 277 "document_condition", 278 [ 279 "pristine exported PDF screenshot", 280 "high-resolution office scanner output", 281 "faded photocopy with mild paper texture", 282 "creased printout with a clipped corner", 283 "low-contrast scan with light shadow near the binding edge", 284 ], 285 ) 286 287 add_category( 288 config_builder, 289 "annotation_layer", 290 [ 291 "no manual annotations", 292 "yellow highlights over two key numbers", 293 "red pen circle around one chart outlier", 294 "blue sticky note partially covering the lower right table", 295 "handwritten margin note asking for follow-up", 296 "rubber stamp reading DRAFT across the header", 297 ], 298 ) 299 300 add_category( 301 config_builder, 302 "numeric_context", 303 [ 304 "include values in thousands with one decimal place", 305 "include percentages, basis-point deltas, and small footnotes", 306 "include dates across the next six months", 307 "include currency values, totals, and year-over-year deltas", 308 "include counts by region plus a total row", 309 ], 310 ) 311 312 config_builder.add_column( 313 dd.ImageColumnConfig( 314 name="document_image", 315 prompt=RICH_DOCUMENT_IMAGE_PROMPT, 316 model_alias=model_alias, 317 ) 318 ) 319 320 return config_builder 321 322 323 def create_dataset( 324 config_builder: dd.DataDesignerConfigBuilder, 325 *, 326 num_records: int, 327 dataset_name: str, 328 artifact_path: Path | str | None = None, 329 ) -> DatasetCreationResults: 330 data_designer = DataDesigner(artifact_path=artifact_path) 331 data_designer.validate(config_builder) 332 return data_designer.create(config_builder, num_records=num_records, dataset_name=dataset_name) 333 334 335 def export_seed_parquet(results: DatasetCreationResults, output_path: Path) -> None: 336 """Export generated images as format-neutral base64 seed rows for VLM pipelines.""" 337 dataset = results.load_dataset() 338 base_path = results.artifact_storage.base_dataset_path 339 rows: list[dict[str, str | int]] = [] 340 341 for row in dataset.itertuples(index=False): 342 image_ref = _first_image_ref(row.document_image) 343 image_path = base_path / image_ref 344 image_format, image_mime_type, image_width, image_height = _image_metadata(image_path) 345 output_row = { 346 "image_base64": base64.b64encode(image_path.read_bytes()).decode("utf-8"), 347 "image_format": image_format, 348 "image_mime_type": image_mime_type, 349 "image_width": image_width, 350 "image_height": image_height, 351 } 352 output_row.update({column: getattr(row, column) for column in SEED_METADATA_COLUMNS}) 353 rows.append(output_row) 354 355 output_path.parent.mkdir(parents=True, exist_ok=True) 356 pd.DataFrame(rows).to_parquet(output_path, index=False) 357 358 359 def _first_image_ref(value: object) -> str: 360 if isinstance(value, str): 361 return value 362 if isinstance(value, Iterable): 363 first = next(iter(value), None) 364 if isinstance(first, str): 365 return first 366 raise ValueError(f"Expected document_image to be a string path or non-empty iterable, got {type(value)!r}") 367 368 369 def _image_metadata(image_path: Path) -> tuple[str, str, int, int]: 370 with Image.open(image_path) as image: 371 image_format = image.format or image_path.suffix.lstrip(".").upper() or "UNKNOWN" 372 image_mime_type = Image.MIME.get(image_format, "application/octet-stream") 373 image_width, image_height = image.size 374 return image_format, image_mime_type, image_width, image_height 375 376 377 RICH_DOCUMENT_IMAGE_PROMPT = """\ 378 Create a realistic single-page business document image with rich visual information. 379 380 Document requirements: 381 - Visual variation ID, for internal diversity only: {{ visual_variation_id }} 382 - Document type: {{ document_type }} 383 - Organization: {{ organization_name }} 384 - Document owner: {{ document_owner }}, {{ owner_role }} 385 - Intended audience: {{ audience }} 386 - Theme: {{ content_theme }} 387 - Layout style: {{ layout_style }} 388 - Physical/rendering condition: {{ document_condition }} 389 - Annotation layer: {{ annotation_layer }} 390 - Numeric style: {{ numeric_context }} 391 392 Required visual content: 393 - Primary visual: {{ primary_visual }} 394 - Secondary visual: {{ secondary_visual }} 395 - At least one readable table with row and column labels 396 - At least one chart, timeline, heatmap, diagram, or KPI-card cluster 397 - A clear title, date, organization name, document owner, section headings, and small source note 398 - Enough readable text to ask visual QA questions about exact values, trends, labels, owners, dates, and relationships 399 400 Make the page visually dense but professionally designed. Use realistic fonts, 401 alignment, legends, axis labels, table borders, captions, and spacing. The text 402 and numbers should be legible. Avoid blank areas, generic placeholder blocks, 403 or lorem ipsum. Generate exactly one final document page for this row. Do not 404 return alternate versions, a grid, a pair of examples, before/after panels, or 405 multiple pages. Use the visual variation ID only as an internal diversity key; 406 never render it as text. Do not include real company logos or real personal 407 data. 408 """ 409 410 411 def parse_args() -> argparse.Namespace: 412 parser = argparse.ArgumentParser(description="Generate rich synthetic business-document images.") 413 parser.add_argument("--num-records", type=int, default=5, help="Number of document images to generate.") 414 parser.add_argument("--dataset-name", default="rich-document-images", help="Output dataset name.") 415 parser.add_argument("--artifact-path", type=Path, default=None, help="Optional Data Designer artifact directory.") 416 parser.add_argument("--model-provider", default=DEFAULT_MODEL_PROVIDER, help="Image model provider name.") 417 parser.add_argument("--model-id", default=DEFAULT_MODEL_ID, help="Provider model ID.") 418 parser.add_argument("--model-alias", default=DEFAULT_MODEL_ALIAS, help="Alias used by image columns.") 419 parser.add_argument("--image-size", default="1K", help="OpenRouter image size tier, such as 1K, 2K, or 4K.") 420 parser.add_argument("--aspect-ratio", default="2:3", help="Provider-specific aspect ratio value.") 421 parser.add_argument("--max-parallel-requests", type=int, default=10, help="Maximum parallel image requests.") 422 parser.add_argument( 423 "--export-seed", 424 type=Path, 425 default=None, 426 help="Optional parquet path for a VQA-ready seed with base64 image bytes and image metadata.", 427 ) 428 return parser.parse_args() 429 430 431 def main() -> None: 432 args = parse_args() 433 config_builder = build_config( 434 model_provider=args.model_provider, 435 model_id=args.model_id, 436 model_alias=args.model_alias, 437 image_size=args.image_size, 438 aspect_ratio=args.aspect_ratio, 439 max_parallel_requests=args.max_parallel_requests, 440 ) 441 results = create_dataset( 442 config_builder, 443 num_records=args.num_records, 444 dataset_name=args.dataset_name, 445 artifact_path=args.artifact_path, 446 ) 447 448 dataset = results.load_dataset() 449 print(f"Generated {len(dataset)} rich document image rows.") 450 print(f"Dataset artifacts: {results.artifact_storage.base_dataset_path}") 451 452 if args.export_seed is not None: 453 export_seed_parquet(results, args.export_seed) 454 print(f"Exported VQA seed parquet: {args.export_seed}") 455 456 457 if __name__ == "__main__": 458 main()