EmoTransfer: Emotion Detecion and Transfer

Emotion Detection and Transfer Project

Stage 1: Emotion Detection Model Fine-Tuning

Data Preparation:

1. Download AffectNet from huggingface and place under emodetection folder.

   code example for downloading dataset from huggingface:

from datasets import load_dataset
dataset = load_dataset('chitradrishti/AffectNet', split='train')  
dataset.to_csv('output.csv') #optional: save dataset as csv file with image paths

3. Run the following command to split the AffectNet into Train and test sets:

cd emodetection
python split_dataset.py

Usage

1. Environment setup: tested with python3.10 + pytorch 2.1.0 + cuda 12.1
2. Download pre-trained EfficientNet checkpoint from here and place under the checkpoints folder.
3. modify any necessary hyper-pameter and run:

bash train.sh

Stage 2: Emotion Transfer with Stable Diffusion Model

Data Preparation:

1. Use test split of Affetnet dataset

Usage

1. Environment setup: tested with python3.10 + pytorch 2.1.0 + cuda 12.1
2. Use stable-diffusion-v1-5 model from huggingface
3. modify any necessary hyper-pameter and run:

python image_generator.py --base_path <YOUR DATA DIR> --emotion_mapping_path <YOUR EMOTION MAPPING FILE> --output_dir <YOUR GENERATED IMAGES DIR>