---
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:

```bash
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:

```bash
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:

```python
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:

```bash
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](https://huggingface.co/docs/hub/spaces-config-reference)