RecipesImage Generation

Rich Document Image Generation

View as Markdown
Download Recipe

Download the complete recipe script

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
13Generate synthetic business-document page images with controlled variation.
14Each generated row pairs an image with the metadata that produced it, making
15the output useful as seed data for visual QA, OCR robustness, multimodal
16judging, and document-understanding experiments.
17
18Prerequisites:
19    - An image-generation provider key for the selected model. The defaults use
20      OpenRouter, so set OPENROUTER_API_KEY before running.
21
22Run:
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
33from __future__ import annotations
34
35import argparse
36import base64
37from collections.abc import Iterable
38from pathlib import Path
39
40import pandas as pd
41from PIL import Image
42
43import data_designer.config as dd
44from data_designer.interface import DataDesigner, DatasetCreationResults
45
46DEFAULT_MODEL_PROVIDER = "openrouter"
47DEFAULT_MODEL_ID = "google/gemini-3.1-flash-image-preview"
48DEFAULT_MODEL_ALIAS = "document-generation-model"
49
50SEED_METADATA_COLUMNS = [
51    "document_type",
52    "primary_visual",
53    "secondary_visual",
54    "layout_style",
55    "document_condition",
56]
57
58
59def 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
89def 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
105def 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
116def 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
323def 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
335def 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
359def _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
369def _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
377RICH_DOCUMENT_IMAGE_PROMPT = """\
378Create a realistic single-page business document image with rich visual information.
379
380Document 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
392Required 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
400Make the page visually dense but professionally designed. Use realistic fonts,
401alignment, legends, axis labels, table borders, captions, and spacing. The text
402and numbers should be legible. Avoid blank areas, generic placeholder blocks,
403or lorem ipsum. Generate exactly one final document page for this row. Do not
404return alternate versions, a grid, a pair of examples, before/after panels, or
405multiple pages. Use the visual variation ID only as an internal diversity key;
406never render it as text. Do not include real company logos or real personal
407data.
408"""
409
410
411def 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
431def 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
457if __name__ == "__main__":
458    main()