AugmentAI 🎨

LLM-Powered Data Augmentation Policy Designer

Design domain-safe, task-aware augmentation policies through natural language conversation. No more manual hyperparameter tuning—just describe your dataset and constraints, and get scientifically sound augmentation pipelines.

Design Philosophy: The LLM suggests. Rules decide. Code executes.

🌟 Features

• One-Command Preparation: augmentai prepare ./dataset - inspect, split, augment, export
• 🆕 AutoSearch: augmentai search ./dataset - find optimal policies automatically
• Natural Language Interface: Describe your dataset in plain English
• Domain-Aware Constraints: Built-in rules for medical, OCR, satellite, and natural images
• Safety-First Design: Hard constraints that prevent scientifically invalid augmentations
• Full Reproducibility: Seed locking, manifest tracking, deterministic pipelines
• LLM as Advisor: OpenAI, Gemini, Ollama, or LM Studio—LLM suggests, rules validate
• Executable Output: Generate standalone Python scripts ready to run
• Verbose/Quiet Modes: --verbose for debugging, --quiet for CI/CD pipelines

🚀 Quick Start

Installation

# Install from PyPI
pip install augmentai

# Or install from source
git clone https://github.com/kyrozepto/augmentai.git
cd augmentai && pip install -e .

Set up your LLM provider

# Option 1: Environment variables (recommended)
export OPENAI_API_KEY="your-api-key"      # For OpenAI
export GOOGLE_API_KEY="your-api-key"       # For Gemini (from aistudio.google.com)

# Option 2: Pass directly via CLI
augmentai chat --provider gemini --api-key "your-key"
augmentai chat --provider openai --api-key "your-key"

# Option 3: Use Ollama (free, local)
ollama pull llama3.2
augmentai chat --provider ollama

Supported LLM Providers:

Provider	Models	API Key Env Var
openai	gpt-4o-mini, gpt-4o	OPENAI_API_KEY
gemini	gemini-2.0-flash, gemini-1.5-pro	GOOGLE_API_KEY
ollama	llama3.2, mistral, etc.	(none needed)
lmstudio	local models	(none needed)

One-Command Dataset Preparation

# Prepare dataset with auto-detected domain
augmentai prepare ./dataset

# Medical domain with custom split
augmentai prepare ./lung_ct --domain medical --split 70/15/15

# Preview what would happen (dry run)
augmentai prepare ./images --dry-run

Interactive Policy Design

# Start interactive chat
augmentai chat --domain medical

# List available domains
augmentai domains

# Validate an existing policy
augmentai validate my_policy.yaml --domain medical

🔍 AutoSearch: Find Optimal Policies Automatically

# Search for best policy (uses evolutionary optimization)
augmentai search ./dataset --domain medical --budget 50

# With custom output directory
augmentai search ./images --budget 100 --output ./search_results

# Preview search configuration
augmentai search ./data --dry-run

AutoSearch uses proxy metrics (diversity, coverage, strength, balance) to score policies without requiring full model training.

� Output Structure

Running augmentai prepare generates:

prepared/
├── data/
│   ├── train/           # Training split
│   ├── val/             # Validation split
│   └── test/            # Test split
├── augmented/           # Output for augmented data
├── augment.py           # Standalone augmentation script
├── config.yaml          # Pipeline configuration
├── manifest.json        # Reproducibility manifest
├── requirements.txt     # Dependencies
└── README.md            # Usage instructions

🎯 Real-World Examples

🏥 Medical: Dental Panoramic X-Ray

augmentai prepare "C:\Users\you\datasets\panoramic-xray" --domain medical --split 70/15/15

# Output:
# ✓ Detected: 1,247 images across 4 classes (cavity, healthy, fracture, caries)
# ✓ Split: 873 train, 187 val, 187 test
# ✓ Policy: Conservative augmentations preserving dental anatomy
# ✓ Forbidden: ElasticTransform, ColorJitter (would distort teeth/bone structure)

🛰️ Satellite: Land Use Classification

augmentai prepare ./sentinel2_tiles --domain satellite --seed 42

# Multi-spectral imagery gets special treatment:
# ✓ Allowed: Rotation (any angle), flips, scale
# ✓ Forbidden: ColorJitter, HSV, ChannelShuffle (breaks spectral bands!)

📝 OCR: Document Scanning

augmentai prepare ./scanned_receipts --domain ocr

# Text legibility is preserved:
# ✓ Allowed: Slight rotation (±5°), brightness adjustment
# ✓ Forbidden: MotionBlur, ElasticTransform (destroys text)

🖼️ Natural: Instagram-style Photos

augmentai prepare ./pet_photos --domain natural

# Maximum flexibility for general images:
# ✓ Allowed: Everything! Color jitter, elastic, cutout, mixup
# ✓ Strong augmentations for robust models

🔬 Research: Reproducible Experiments

# Exact same augmentations every time
augmentai prepare ./experiment_data --seed 12345 --output ./exp_v1

# Later, reproduce with the manifest
cat ./exp_v1/manifest.json
# {
#   "seed": 12345,
#   "dataset_hash": "a1b2c3d4...",
#   "policy_hash": "e5f6g7h8...",
#   "augmentai_version": "0.1.0"
# }

🏭 Production: Batch Processing

# Prepare multiple datasets with consistent settings
for dataset in chest_xray brain_mri skin_lesion; do
    augmentai prepare ./raw/$dataset --domain medical --output ./prepared/$dataset
done

💻 Python API: Full Control

from augmentai.core import Policy, Transform
from augmentai.domains import get_domain
from augmentai.rules.enforcement import RuleEnforcer
from augmentai.export import ScriptGenerator

# Create custom policy
policy = Policy(
    name="dental_xray_v2",
    domain="medical",
    transforms=[
        Transform("HorizontalFlip", 0.5),
        Transform("Rotate", 0.7, parameters={"limit": 20}),
        Transform("CLAHE", 0.4, parameters={"clip_limit": 2.0}),
        Transform("GaussNoise", 0.3, parameters={"var_limit": (10, 30)}),
        Transform("RandomBrightnessContrast", 0.5),
    ]
)

# Validate against medical rules
enforcer = RuleEnforcer(get_domain("medical"))
result = enforcer.enforce_policy(policy)
print(f"Safe transforms: {len(result.policy.transforms)}")

# Generate standalone script
generator = ScriptGenerator()
script = generator.generate_augment_script(result.policy, seed=42)
with open("my_augment.py", "w") as f:
    f.write(script)

🔄 Modifying Generated Config

After running augmentai prepare, edit config.yaml:

# config.yaml - tweak as needed!
policy:
  name: dental_xray_policy
  domain: medical
  seed: 42
transforms:
  - name: HorizontalFlip
    probability: 0.5          # ← Increase to 0.7?
  - name: Rotate
    probability: 0.5
    parameters:
      limit: 15               # ← Bump to 25 for more variety
  - name: CLAHE               # ← Add new transforms!
    probability: 0.4
    parameters:
      clip_limit: 2.0

Then re-run the generated script:

python augment.py --input data/train --output augmented/train --seed 42

🎓 Interactive Learning Session

$ augmentai chat --domain medical

You: I have dental panoramic x-rays with cavity annotations. What augmentations 
     are safe? I'm worried about distorting the teeth.

🤖 AugmentAI: Great concern! For dental x-rays, I recommend:

✅ Safe for dental imaging:
  • HorizontalFlip (p=0.5) - Teeth are roughly symmetric
  • Rotate (p=0.5, limit=±10°) - Small rotations are realistic
  • CLAHE (p=0.3) - Enhances contrast in x-rays
  • GaussNoise (p=0.2) - Simulates sensor noise

⚠️ Automatically blocked:
  • ElasticTransform - Would warp tooth boundaries unrealistically
  • GridDistortion - Same issue, breaks tooth geometry  
  • ColorJitter - X-rays are grayscale, color transforms invalid

You: Can I increase rotation to 30 degrees?

🤖 AugmentAI: I'd advise caution. Dental panoramic x-rays are typically taken 
with the patient positioned consistently. Rotations beyond ±15° might create 
unrealistic training samples. However, if your dataset has varied positioning, 
±20° could work. I'll set it to 20° as a compromise.

/preview
┌────────────────────────────┬─────────────┬──────────────────────┐
│ Transform                  │ Probability │ Parameters           │
├────────────────────────────┼─────────────┼──────────────────────┤
│ HorizontalFlip             │ 50%         │ -                    │
│ Rotate                     │ 50%         │ limit=20             │
│ CLAHE                      │ 30%         │ clip_limit=2.0       │
│ GaussNoise                 │ 20%         │ var_limit=(5, 25)    │
│ RandomBrightnessContrast   │ 30%         │ brightness_limit=0.1 │
└────────────────────────────┴─────────────┴──────────────────────┘

/export dental_cavity_policy.py
✓ Exported to dental_cavity_policy.py

🏥 Domain Safety

AugmentAI enforces hard constraints that cannot be overridden:

Domain	Forbidden Transforms	Reason
Medical	ElasticTransform, GridDistortion, ColorJitter	Breaks anatomy, invalid for grayscale
OCR	MotionBlur, ElasticTransform	Destroys text legibility
Satellite	ColorJitter, HSV, ChannelShuffle	Breaks spectral band relationships

📁 Project Structure

augmentai/
├── cli/              # CLI commands (prepare, chat, validate, search)
├── core/             # Policy, Transform, Manifest, Pipeline
├── domains/          # Domain rules (medical, ocr, satellite, natural)
├── inspection/       # Dataset auto-detection & analysis
├── splitting/        # Train/val/test splitting strategies
├── export/           # Script & folder generation
├── llm/              # LLM client and prompts
├── rules/            # Safety validation & enforcement
├── compilers/        # Backend code generation
├── search/           # 🆕 AutoSearch: evolutionary policy optimization
├── utils/            # 🆕 Progress bars, logging utilities
└── exceptions.py     # 🆕 Custom error hierarchy

🔧 Configuration

Create augmentai.yaml in your project:

llm:
  provider: openai  # or ollama, lmstudio
  model: gpt-4o-mini
  temperature: 0.7

backend: albumentations
output_dir: ./policies

🎯 Custom Domains

Define your own domain constraints in YAML:

# my_domain.yaml
name: my_custom_domain
description: Custom constraints for my task

constraints:
  - transform_name: ElasticTransform
    level: forbidden
    reason: Not suitable for my task

  - transform_name: Rotate
    level: recommended
    reason: Good for rotation invariance
    parameter_limits:
      limit: [-30, 30]

recommended_transforms:
  - HorizontalFlip
  - RandomBrightnessContrast

Load with:

augmentai chat --domain-file my_domain.yaml

🤝 Contributing

Contributions welcome! See CONTRIBUTING.md for guidelines.

📄 License

MIT License - see LICENSE for details.

🙏 Acknowledgments

• Albumentations for the augmentation backend
• Rich for beautiful terminal UI
• OpenAI, Ollama, LM Studio for LLM support