README.md

metadata

title: ColorGradingAE
emoji: 🎨
colorFrom: purple
colorTo: indigo
sdk: gradio
sdk_version: 6.20.0
python_version: '3.13'
app_file: gradio_demo2.py
pinned: false
license: mit

ColorGradingAE - Contrastive Color Grade AI

An AI-powered color grading tool that uses contrastive learning with an encoder-decoder architecture to automatically analyze and apply professional color grading to images.

Features

• 🎨 Automatic Color Analysis: Extracts color features from images using HSV color space analysis
• 🧠 Contrastive Learning Model: Uses encoder-decoder architecture trained on color grading patterns
• 🎯 Zone-based Processing: Analyzes shadows, midtones, and highlights separately
• 🌐 Interactive Web Interface: Easy-to-use Gradio interface for real-time color grading
• 🚀 GPU Acceleration: CUDA support for fast inference

Installation

Clone the repository and install dependencies:

pip install -r requirements.txt

Requirements

• Python 3.13+
• PyTorch with CUDA support (or CPU)
• OpenCV
• Gradio
• Pillow
• NumPy

Usage

Web Interface

Run the interactive Gradio demo:

python gradio_demo.py

This launches a web interface where you can upload images and apply color grading in real-time.

Python API

Use the color grading functions programmatically:

from inference import generate_grade, apply_grade

# Generate color grade parameters for an image
grade = generate_grade("input.jpg")

# Apply the generated grade to an image
output = apply_grade("input.jpg", grade)

Training

To train or fine-tune the model on your own images:

python trainer.py

Place your training images in the data/images/ directory.

Project Structure

├── gradio_demo.py          # Main Gradio web interface
├── gradio_demo2.py         # Alternative demo interface
├── inference.py            # Color grading inference functions
├── trainer.py              # Model training and feature extraction
├── downloader.py           # Utility for downloading images/models
├── encoder.pt              # Pre-trained encoder model
├── decoder.pt              # Pre-trained decoder model
├── data/
│   └── images/             # Training images directory
└── flagged/                # Gradio flagged/saved results

How It Works

1. Feature Extraction: Analyzes image color distributions across HSV channels in 16 hue bins
2. Zone Analysis: Processes shadows, midtones, and highlights separately for more nuanced grading
3. Encoding: Compresses color features using the trained encoder
4. Decoding: Generates color grading parameters using the trained decoder
5. Application: Applies the generated grade to produce the final color-graded output

Model Architecture

• Encoder: Compresses color feature vectors into latent space
• Decoder: Reconstructs color grading parameters from latent vectors
• Training: Uses contrastive learning to learn meaningful color relationships

Performance

• Supports GPU acceleration via CUDA for fast inference
• Falls back to CPU if GPU unavailable
• Real-time processing suitable for interactive applications

License

MIT License - See LICENSE file for details

---

For more information, check out the Hugging Face documentation