electblake/image-data-extractor

1import zipfile2from pathlib import Path3 4import gradio as gr5 6from model import extract, generate_template7 8templates_directory = Path(__file__).parent / "data" / "templates"9structured_json_templates = {10    "Basic receipt": (templates_directory / "basic-receipt.json").read_text(11        encoding="utf-8"12    ),13    "Invoice with line items": (14        templates_directory / "invoice-with-line-items.json"15    ).read_text(encoding="utf-8"),16    "Basic bank statement": (17        templates_directory / "basic-bank-statement.json"18    ).read_text(encoding="utf-8"),19    "Advanced bank statement": (20        templates_directory / "advanced-bank-statement.json"21    ).read_text(encoding="utf-8"),22    "Model task catalog": (23        templates_directory / "model-task-catalog.json"24    ).read_text(encoding="utf-8"),25    "Todo list": (templates_directory / "todo-list.json").read_text(26        encoding="utf-8"27    ),28}29 30default_structured_json_template = "Invoice with line items"31 32 33def download_agentskill(include_all_skills):34    skill_folders = (35        sorted(path for path in Path(".agents/skills").iterdir() if path.is_dir())36        if include_all_skills37        else [Path(".agents/skills/image-data-extractor")]38    )39    for skill_folder in skill_folders:40        with zipfile.ZipFile(41            f"{skill_folder.name}.zip", "w", compression=zipfile.ZIP_DEFLATED42        ) as archive:43            for path in skill_folder.rglob("*"):44                archive.write(path, path.relative_to(".agents/skills"))45    return [f"{skill_folder.name}.zip" for skill_folder in skill_folders]46 47with gr.Blocks(title="Image Data Extractor") as demo:48    gr.Markdown(49        "# Image Data Extractor\n"50        "Extract structured JSON from an image and optional text using "51        "[numind/NuExtract3](https://huggingface.co/numind/NuExtract3)."52    )53    with gr.Row():54        with gr.Column():55            image = gr.Image(label="Document image", type="pil", height=430)56            text = gr.Textbox(57                label="Document text (optional)",58                placeholder="Add text to process alongside the image",59                lines=3,60            )61            enable_thinking = gr.Checkbox(label="Enable thinking")62            template_preset = gr.Dropdown(63                choices=list(structured_json_templates),64                value=default_structured_json_template,65                label="JSON template preset",66                info="Choose a starting structure, then edit it below.",67            )68            generate_template_button = gr.Button("Generate template from Image")69            with gr.Accordion(70                "Structured JSON template", open=False71            ) as template_accordion:72                template = gr.Textbox(73                    label="Structured JSON template",74                    value=structured_json_templates[default_structured_json_template],75                    lines=18,76                    info="Used for structured extraction. Edit the example to match your document.",77                )78            run = gr.Button("Extract Image Data", variant="primary")79        with gr.Column():80            structured_output = gr.JSON(81                label="Structured JSON",82                open=True,83                show_indices=True,84            )85            gr.Markdown(86                "## Install the agent skill\n"87                "Download the ZIP, extract it, and place its skill folders in your "88                "project's `.agents/skills/` directory. Agents should read the "89                "`SKILL.md` and use it whenever a task requires structured extraction "90                "from an image. Keep **Include All Skills** checked to download every "91                "skill as its own ZIP file."92            )93            include_all_skills = gr.Checkbox(94                label="Include All Skills",95                value=True,96            )97            download_agentskill_button = gr.Button("Download agent skill")98            download_agentskill_file = gr.File(99                label="Agent skill ZIPs",100                file_count="multiple",101            )102 103    gr.Examples(104        examples=[105            [106                preset,107                str(Path(__file__).parent / "data" / "samples" / filename),108                "",109                False,110            ]111            for preset, filename in [112                ("Basic receipt", "receipt-ocr-original.webp"),113                ("Invoice with line items", "invoice-with-items.png"),114                ("Advanced bank statement", "BankStatementChequing.png"),115                ("Basic bank statement", "bank statement blog image.webp"),116                ("Model task catalog", "tags.png"),117                ("Todo list", "task-list.png"),118            ]119        ],120        inputs=[template_preset, image, text, enable_thinking],121    )122 123    template_preset.change(124        structured_json_templates.__getitem__,125        inputs=template_preset,126        outputs=template,127        api_name="select_structured_json_template",128    )129 130    download_agentskill_button.click(131        download_agentskill,132        inputs=include_all_skills,133        outputs=download_agentskill_file,134        api_name="download_agentskill",135    )136 137    generate_template_button.click(138        lambda: gr.Accordion(open=True),139        outputs=template_accordion,140        api_visibility="private",141        show_progress="hidden",142    ).then(143        generate_template,144        inputs=image,145        outputs=template,146        api_name="generate_template",147    )148 149    run.click(150        extract,151        inputs=[image, text, template, enable_thinking],152        outputs=structured_output,153        api_name="extract",154    )155 156 157if __name__ == "__main__":158    demo.launch()159