<div align="center">

<img src="img/logo.png" alt="AutoFigure Logo" width="100%"/>


# AutoFigure: Generating and Refining Publication-Ready Scientific Illustrations [ICLR 2026]

[![ICLR 2026](https://img.shields.io/badge/ICLR-2026-blue?style=for-the-badge&logo=openreview)](https://openreview.net/forum?id=5N3z9JQJKq)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg?style=for-the-badge)](https://opensource.org/licenses/MIT)
[![Python](https://img.shields.io/badge/Python-3.8%2B-blue?style=for-the-badge&logo=python&logoColor=white)](https://www.python.org/)
[![HuggingFace](https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-FigureBench-orange?style=for-the-badge)](https://huggingface.co/datasets/WestlakeNLP/FigureBench)

<p align="center">
  <strong>From Text to Publication-Ready Diagrams</strong><br>
  AutoFigure is an intelligent system that leverages Large Language Models (LLMs) with iterative refinement to generate high-quality scientific figures from text descriptions or research papers.
</p>

[Quick Start](#-quick-start) • [Web Interface](#-web-interface) • [Configuration](#%EF%B8%8F-configuration) • [API Reference](#-api-reference)

</div>

---

## ✨ Features

| Feature | Description |
| :--- | :--- |
| 📝 **Text-to-Figure** | Generate figures directly from natural language descriptions. |
| 📄 **Paper-to-Figure** | Extract methodology from PDFs and create visual diagrams automatically. |
| 🔄 **Iterative Refinement** | Dual-agent system (Generation + Evaluation) for continuous quality optimization. |
| 🎨 **Multiple Formats** | Output as **SVG** or **mxGraph XML** (fully compatible with draw.io). |
| 💅 **Image Enhancement** | Optional AI-powered post-processing for aesthetic beautification. |
| 🖥️ **Web Interface** | Interactive Next.js frontend for easy generation and editing. |

---

## 🚀 How It Works

AutoFigure employs a **Review-Refine** loop to ensure high accuracy and aesthetic quality.

<div align="center">
<img src="img/method.png" alt="AutoFigure method" width="1000"/>
</div>

> **Process:**
> 1. **Generate:** The agent creates initial SVG/XML based on description & references.
> 2. **Evaluate:** The critic scores quality (0-10) and provides specific feedback.
> 3. **Refine:** The loop continues until the figure meets publication standards.

---

## 🌟 Generated Examples

Here are examples of figures generated by AutoFigure across different domains, showcasing its versatility in handling various levels of complexity.

| Category & Visualization |
| :---: |
| **📄 Paper Case**<br><img src="img/case_paper.png" width="100%" alt="Paper Case"/> |
| **📊 Survey Case**<br><img src="img/case_survey.png" width="100%" alt="Survey Case"/> |
| **📝 Blog Case**<br><img src="img/case_blog.png" width="100%" alt="Blog Case"/> |
| **📘 Textbook Case**<br><img src="img/case_textbook.png" width="100%" alt="Textbook Case"/> |

---

## ⚡ Quick Start

### Option 1: Python SDK (Recommended)


You can install via cloning the repo:

```bash
git clone https://github.com/ResearAI/AutoFigure.git
cd AutoFigure
pip install -e .
playwright install chromium  # Required for rendering
```

#### 1. Basic Usage (Text-to-Figure)

```python
from autofigure import AutoFigureAgent, Config

# 1. Configure
config = Config(
    generation_api_key="your-api-key",
    generation_provider="openrouter",  # options: 'openrouter', 'gemini', 'bianxie'
    generation_model="google/gemini-2.5-pro",
)

# 2. Generate
agent = AutoFigureAgent(config)
result = agent.generate(
    description="A flowchart showing transformer training pipeline",
    max_iterations=5,
    output_format="svg",
    topic="paper" # 'paper', 'survey', 'blog', 'textbook'
)

print(f"✅ Generated: {result.svg_path} (Score: {result.final_score}/10)")
```

#### 2. Generate from Paper (PDF/Markdown)

Extract methodology from a paper and generate a figure automatically.

```python
# Generate figure from paper (PDF or Markdown)
result = agent.generate_from_paper(
    paper_path="./paper.pdf",
    max_iterations=5,
    output_format="svg",
    enable_enhancement=True, # Enhance the result
)

if result.success:
    print(f"Extracted methodology: {result.methodology_text[:200]}...")
    print(f"Generated figure: {result.svg_path}")
```

#### 3. With Image Enhancement

Generate multiple enhanced aesthetic variants of the figure.

```python
result = agent.generate(
    description="Neural network architecture diagram",
    enable_enhancement=True,
    enhancement_count=3,     # Generate 3 variants
    art_style="Modern scientific illustration with clean lines",
    enhancement_input_type="code2prompt" # Best quality mode
)

if result.success:
    print(f"Original Preview: {result.preview_path}")
    print(f"Enhanced variants: {result.enhanced_paths}")
```

### Option 2: Web Interface

Ideally suited for visual interaction and editing.

```bash
./start.sh
# Then open http://localhost:6002 in your browser
```

---

## 📊 FigureBench Dataset

We introduce **FigureBench**, the first large-scale benchmark for generating scientific illustrations from long-form text.

<div align="center">
<img src="img/figurebench.png" alt="figurebench" width="1000"/>
</div>

### Dataset Overview

| Category | Samples | Avg. Tokens | Text Density | Complexity |
|:---|:---:|:---:|:---:|:---:|
| 📄 **Paper** | 3,200 | 12,732 | 42.1% | High |
| 📝 **Blog** | 20 | 4,047 | 46.0% | Med |
| 📊 **Survey** | 40 | 2,179 | 43.8% | High |
| 📘 **Textbook** | 40 | 352 | 25.0% | Low |
| **Total** | **3,300** | **10k+** | **41.2%** | **~5.3 Components** |

### Download
<div align="left">
  <a href="https://huggingface.co/datasets/WestlakeNLP/FigureBench">
    <img src="https://img.shields.io/badge/%F0%9F%A4%97%20Download%20on%20HuggingFace-FFD21E?style=for-the-badge&logoColor=black" alt="Download">
  </a>
</div>

```python
from datasets import load_dataset
dataset = load_dataset("WestlakeNLP/FigureBench")
```

---

## ⚙️ Configuration

AutoFigure is highly configurable. You can set these in `Config()` or via environment variables.

### Supported LLM Providers

| Provider | Base URL | Recommended Models |
|----------|----------|--------------------|
| **OpenRouter** | `openrouter.ai/api/v1` | `gemini-2.5-pro` |
| **Bianxie** | `api.bianxie.ai/v1` | `gemini-2.5-pro` |
| **Google** | `generativelanguage...` | `gemini-2.5-pro` |

### Generation Settings

| Option | Description | Default |
|--------|-------------|---------|
| `generation_api_key` | API key for figure generation | Required |
| `generation_base_url` | Base URL for API | Provider default |
| `generation_model` | Model name | Provider default |
| `generation_provider` | Provider: 'openrouter', 'bianxie', 'gemini' | 'openrouter' |

### Methodology Extraction Settings

| Option | Description | Default |
|--------|-------------|---------|
| `methodology_api_key` | API key for methodology extraction | Same as generation |
| `methodology_model` | Model for methodology extraction | Same as generation |
| `methodology_provider` | Provider for methodology extraction | Same as generation |

### Enhancement Settings

| Option | Description | Default |
|--------|-------------|---------|
| `enhancement_api_key` | API key for image enhancement | None |
| `enhancement_provider` | Enhancement provider | 'openrouter' |
| `enhancement_model` | Model for image enhancement | Provider default |
| `enhancement_input_type` | Input type: 'none', 'code', 'code2prompt' | 'code2prompt' |
| `enhancement_count` | Number of enhanced variants to generate | 1 |
| `art_style` | Art style description for enhancement | '' |

### Pipeline Settings

| Option | Description | Default |
|--------|-------------|---------|
| `max_iterations` | Maximum refinement iterations | 5 |
| `quality_threshold` | Quality threshold (0-10) | 9.0 |
| `output_dir` | Output directory | './autofigure_output' |
| `custom_references` | Custom reference figure paths | None |

---

## 📚 API Reference

### `generate()` Parameters

| Parameter | Description |
|-----------|-------------|
| `description` | Text description of the figure to generate |
| `max_iterations` | Maximum iterations (overrides config) |
| `output_format` | 'svg' or 'mxgraphxml' |
| `quality_threshold` | Quality threshold (overrides config) |
| `enable_enhancement` | Whether to enhance the final image |
| `art_style` | Art style for enhancement (overrides config) |
| `enhancement_input_type` | 'none', 'code', or 'code2prompt' (overrides config) |
| `enhancement_count` | Number of enhanced variants (overrides config) |
| `topic` | Content type: 'paper', 'survey', 'blog', 'textbook' |
| `custom_references` | Custom reference figure paths |

### `generate_from_paper()` Parameters

Accepts all parameters from `generate()` plus:

| Parameter | Description |
|-----------|-------------|
| `paper_path` | Path to paper file (PDF or Markdown) |
| `methodology_api_key` | API key for extraction (overrides config) |
| `methodology_provider` | Provider for extraction (overrides config) |

### Result Object (`GenerationResult`)

| Attribute | Description |
|-----------|-------------|
| `success` | Whether generation was successful |
| `svg_path` | Path to generated SVG file |
| `mxgraph_path` | Path to generated mxGraph XML file |
| `preview_path` | Path to PNG preview image |
| `enhanced_paths` | List of all enhanced image paths |
| `final_score` | Final quality score (0-10) |
| `methodology_text` | Extracted methodology (from paper) |
| `error` | Error message if failed |

### Enhancement Modes

| Mode | Description |
|------|-------------|
| `none` | Direct beautification without code reference |
| `code` | Use generated code (SVG/XML) as reference |
| `code2prompt` | Use LLM to analyze code and generate detailed prompt (recommended) |

---

## 📁 Project Structure

<details>
<summary>Click to expand directory tree</summary>

```
AutoFigure/
├── autofigure/              # 📦 Python SDK
│   ├── agent.py             # Main Agent
│   ├── generator.py         # Generation Pipeline
│   ├── enhancer.py          # Image Enhancement
│   └── extractor.py         # PDF Method Extraction
├── frontend/                # 🖥️ Next.js Web UI
├── backend/                 # 🔌 Flask API Server
├── scripts/                 # 🛠️ Utility Scripts
└── pyproject.toml           # Config
```
</details>

---

## 🤝 Community & Support

**WeChat Discussion Group**  
Scan the QR code to join our community. If the code is expired, please add WeChat ID `nauhcutnil` or contact `tuchuan@mail.hfut.edu.cn`.
<table>
  <tr>
    <td><img src="img/wechat3.jpg" width="200" alt="WeChat 1"/></td>
    <td><img src="img/wechat4.jpg" width="200" alt="WeChat 2"/></td>
  </tr>
</table>
---

## 📜 Citation & License

If you use **AutoFigure** or **FigureBench** in your research, please cite:

```bibtex
@inproceedings{
zhu2026autofigure,
title={AutoFigure: Generating and Refining Publication-Ready Scientific Illustrations},
author={Minjun Zhu and Zhen Lin and Yixuan Weng and Panzhong Lu and Qiujie Xie and Yifan Wei and Sifan Liu and Qiyao Sun and Yue Zhang},
booktitle={The Fourteenth International Conference on Learning Representations},
year={2026},
url={https://openreview.net/forum?id=5N3z9JQJKq}
}
```
