davi-rvq 0.1.7
RVQ-VAE for motion capture data compression and reconstruction
License: MIT
Requires Python: >=3.8
Keywords: machine-learning, vqvae, rvqvae, motion-capture, compression

# DAVI RVQ-VAE

**Residual Vector Quantized Variational AutoEncoder for Motion Capture Data Compression**

A high-performance PyTorch library for compressing and reconstructing motion capture data using Residual Vector Quantization (RVQ). Achieve up to **4192x compression** with median angular error of only **5.17°**.

## Features

- **High Compression Ratio**: Up to 4192x compression of motion capture data
- **Low Reconstruction Error**: Median MAE of 5.17° (angular)
- **Multiple Encoder Types**: Flatten, Conv1D, Deep Conv1D, Attention
- **Flexible Architecture**: Support for VQ-VAE and RVQ-VAE
- **Easy to Use**: Simple CLI for training, inference, and demo
- **Pre-trained Model**: Includes MAXIMUM_QUALITY_STRIDE4 model

## Installation

### From PyPI (once published)

```bash
pip install davi_rvq
```

### From Source

```bash
git clone https://github.com/yourusername/davi_rvq.git
cd davi_rvq
pip install -e .
```

## Quick Start

### 1. Training

Train a new model with your own data:

```bash
davi-train --config examples/configs/default_config.yaml
```

### 2. Inference

Run inference on CSV files:

```bash
davi-infer --checkpoint davi_rvq/checkpoints/model_best.pth \
           --input data.csv \
           --output reconstruction.csv
```

### 3. Demo & Showcase

Create professional showcases with skeleton animations and video comparisons:

```bash
# Animation only (uses default model if available)
davi-demo --csv data.csv --output-dir output/

# With original video (side-by-side)
davi-demo --csv data.csv --video original.mp4 --output-dir output/

# With 3 videos (original | animation | comparison)
davi-demo --csv data.csv --video original.mp4 \
          --comparison comparison.mp4 --output-dir output/
```

**Output files:**
- `reconstruction_skeleton_4views.mp4` - 4-panel skeleton animation
- `comparison_2panel.mp4` or `comparison_3panel.mp4` - Final comparison
- `reconstruction_angles.csv` - Reconstructed motion data
- `rvq_codes.npy` - Discrete compression codes (with `--save-codes`)

**Note**: `davi-demo` creates full visual showcases. For simple CSV reconstruction, use `davi-infer` instead.

See [CLI_USAGE.md](CLI_USAGE.md) for detailed command documentation.

## Python API

### Load Pre-trained Model

```python
import torch
from davi_rvq.inference import load_model, reconstruct_from_csv

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Load model
model, config, norm_stats = load_model('davi_rvq/checkpoints/model_best.pth', device)

# Reconstruct from CSV
original, reconstructed = reconstruct_from_csv(
    model, 'input.csv', norm_stats, device,
    window_size=config['data']['window_size'],
    stride=4
)
```

### Encode and Decode

```python
import numpy as np
from davi_rvq.inference import encode_window, decode_codes

# Prepare window (shape: [window_size, num_features])
window = np.random.randn(16, 144)  # Example: 16 frames, 144 features

# Encode to discrete codes
codes = encode_window(model, window, norm_stats, device)
print(f"Codes shape: {codes.shape}")  # (1, num_codebooks)

# Decode back to motion
reconstruction = decode_codes(model, codes, norm_stats, device)
print(f"Reconstruction shape: {reconstruction.shape}")  # (1, num_features)
```

### Training Custom Model

```python
from davi_rvq import RVQVAE, SlidingWindowDataset, Trainer, find_csv_files
from torch.utils.data import DataLoader
import torch

# Create dataset
csv_files = find_csv_files('data/', pattern='*.csv')
dataset = SlidingWindowDataset(csv_files, window_size=16, stride=4)

# Create model
model = RVQVAE(
    input_dim=144,
    window_size=16,
    hidden_dim=1024,
    z_dim=512,
    num_embeddings=2048,
    num_codebooks=12,
    commitment_cost=0.25,
    encoder_type='deep_conv1d',
    decoder_type='residual'
)

# Train
device = torch.device('cuda')
model = model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

train_loader = DataLoader(dataset, batch_size=256, shuffle=True)
# ... (see examples/ for full training script)
```

## Model Architecture

The library supports two main architectures:

### VQ-VAE (Vector Quantized VAE)
- Single codebook quantization
- Fast encoding/decoding
- Good for basic compression

### RVQ-VAE (Residual Vector Quantized VAE)
- Multiple codebook layers (typically 4-12)
- Hierarchical residual quantization
- Better reconstruction quality
- **Recommended for production use**

### Encoder Types
- `flatten`: Simple MLP encoder
- `conv1d`: 1D Convolutional encoder
- `deep_conv1d`: Deep convolutional with residual connections
- `attention`: Transformer-based attention encoder

### Decoder Types
- `standard`: Simple MLP decoder
- `deep`: Deep MLP with more layers
- `residual`: Residual connections for better gradient flow

## Pre-trained Model

The package includes a pre-trained model (`model_best.pth`) with the following configuration:

- **Architecture**: RVQ-VAE with Deep Conv1D encoder and Residual decoder
- **Window Size**: 16 frames
- **Stride**: 4 frames
- **Codebooks**: 12
- **Embedding Dimension**: ~43 per codebook
- **Total Embeddings**: 2048 per codebook
- **Compression Ratio**: 4192x
- **Performance**:
  - Median MAE: 5.17°
  - 47% of samples < 5° error

## Configuration

Example configuration file (`examples/configs/default_config.yaml`):

```yaml
experiment:
  name: "rvqvae_MAXIMUM_QUALITY_STRIDE4"
  seed: 42

model:
  mode: "rvq"  # or "vq"
  input_dim: null  # auto-detected
  hidden_dim: 1024
  z_dim: 512
  num_embeddings: 2048
  num_codebooks: 12
  commitment_cost: 0.25
  encoder_type: "deep_conv1d"
  decoder_type: "residual"

data:
  data_dir: "data/"
  pattern: "*.csv"
  window_size: 16
  stride: 4
  frame_column: "frame"
  train_split: 0.8
  val_split: 0.15

training:
  batch_size: 256
  lr: 0.0001
  num_epochs: 100
  # ... (see examples/configs/ for full config)
```

## Examples

See the `examples/` directory for:

- **Training examples**: Different configurations for various use cases
- **Showcase demos**: Visualization scripts and demo videos
- **Evaluation scripts**: Tools to evaluate model performance

## Project Structure

```
davi_rvq/
├── davi_rvq/              # Main package
│   ├── models/            # Model architectures
│   ├── data/              # Dataset utilities
│   ├── utils/             # Training utilities
│   ├── inference/         # Inference utilities
│   ├── cli/               # Command-line interfaces
│   └── checkpoints/       # Pre-trained models
├── examples/              # Examples and demos
│   ├── configs/           # Configuration files
│   ├── showcase/          # Demo videos and visualizations
│   └── scripts/           # Example scripts
├── setup.py               # Installation script
├── pyproject.toml         # Package metadata
└── README.md              # This file
```

## Requirements

- Python >= 3.8
- PyTorch >= 2.0.0
- NumPy >= 1.24.0
- Pandas >= 2.0.0
- PyYAML >= 6.0
- (Optional) Weights & Biases for experiment tracking

## Citation

If you use this library in your research, please cite:

```bibtex
@software{davi_rvqvae,
  title={DAVI RVQ-VAE: Motion Capture Compression with Residual Vector Quantization},
  author={Davi},
  year={2025},
  url={https://github.com/yourusername/davi_rvq}
}
```

## License

MIT License - see LICENSE file for details

## Contributing

Contributions are welcome! Please feel free to submit a Pull Request.

## Support

For issues and questions:
- Open an issue on GitHub
- Check the `examples/showcase/` directory for demos and documentation

## Changelog

### v0.1.0 (2025-01-23)
- Initial release
- Pre-trained MAXIMUM_QUALITY_STRIDE4 model
- CLI tools for train, infer, and demo
- Comprehensive documentation and examples