# ExploreAug: Exploratory Data Augmentation Framework

An advanced data augmentation framework combining diffusion models and dataset cartography for improving deep learning model performance and robustness.

## Requirements

To install the requirements, run the following in your environment first:
```bash
pip install torch torchvision torchaudio
pip install diffusers transformers timm
pip install omegaconf tensorboard
pip install matplotlib pillow scikit-learn
pip install psutil opencv-python
```

You can also download the pretrained Stable Diffusion model from [Huggingface](https://huggingface.co/stabilityai/stable-diffusion-2-1-base), and the code will automatically fallback to online download if local model is not available.

## Dataset Setup

### Dataset Structure
```
data/
├── afhq/
│   ├── train/
│   │   ├── cat/
│   │   ├── dog/
│   │   └── wild/
│   └── val/
└── ...
```

## Run the Code

You can run the code with the following command:
```bash
python main.py --config configs/exploreaug/resnet-afhq.yaml
```

The script supports the following options:
- `--config`: Path to the configuration file (required)


