Improvement to the Vision Model - include wider range of emotions

[notzabir]

Current Vision System Overview

The BrainCharge project has a computer vision pipeline with the following architecture:

How It Works:

1. Face Detection: YOLOv8 (yolov8n-face-lindevs.pt) detects faces in real-time video
2. Emotion Classification: EfficientNet-B2 model (emotions_model.pt) classifies detected faces into emotions
3. Person Tracking: YOLOv8n tracks people for robot movement commands
4. Image Processing: Faces are preprocessed (288x288 grayscale→RGB, normalized)

Current Features:

Feature	Implementation
Supported Emotions	5 classes: Angry, Fear, Happy, Neutral, Sad
Model Architecture	EfficientNet-B2 with custom classifier head
Training Dataset	FER2013 (excludes Disgust and Surprise)
Top Predictions	Returns top 2 emotions with confidence >0.1
Camera Integration	Multiple backend support (DirectShow, ALSA, PulseAudio)
Movement Tracking	ByteTrack for person tracking + Arduino control
Data Augmentation	Horizontal flip, rotation (±20°), color jitter

---

Improvement Plan: Wider Range of Emotions

Phase 1: Expand Emotion Classes

Goal: Include Disgust and Surprise + 2-3 additional emotions

Resources Needed:

1. Dataset Expansion

   • Keep FER2013 base but include excluded classes (Disgust, Surprise)
   • Add supplementary datasets: AffectNet, CK+, JAFFE
   • Total target: 8-10 emotion classes

2. Data Collection & Preprocessing

   • Standardize images to 288x288 with balanced class distribution
   • Implement class weighting to handle imbalance (already partially done in train_classifier.py)
   • Update fer2013.py to include all 7 classes:

     AVAILABLE_EMOTIONS = ["Angry", "Disgust", "Fear", "Happy", "Neutral", "Sad", "Surprise"]

3. Model Architecture Updates

   • Update cv_model.py: Change final layer output from 5 to 7 neurons
   • Adjust picture.py: Line 22 and Line 74 (output layer)
   • Consider ensemble approach: combine EfficientNet-B2 with ResNet50 for better accuracy

Phase 2: Model Retraining

Resources:

• GPU compute (recommended: NVIDIA A100 or RTX 4090 for faster training)
• Training time: ~6-8 hours for 60 epochs on expanded dataset
• Memory: ~16GB VRAM minimum

Steps:

# Update constants in train_classifier.py
NUM_EPOCHS = 100  # Increase iterations
class_weights = [higher_weight_for_rare_classes]  # Balance dataset

Phase 3: Fine-tuning & Validation

• Add K-fold cross-validation to train_classifier.py
• Implement confidence thresholding: adjust line 336 in picture.py from 0.1 to 0.15-0.20
• Test on real-world caregiver scenarios

---

Detailed Implementation Roadmap

Code Changes Required:

1. cv/cv_model.py - Change output classes:

   nn.Linear(512, 7, bias=True)  # Was: 5

2. cv/picture.py - Update constants:

   AVAILABLE_EMOTIONS = ["Angry", "Disgust", "Fear", "Happy", "Neutral", "Sad", "Surprise"]
   NUM_TOP_EMOTIONS = 3  # Can show top 3 instead of 2

3. cv/fer2013.py - Remove emotion filtering:

   # Remove this line:
   # if d.name not in ["disgust", "surprise"]
   # Keep all 7 emotions

4. cv/train_classifier.py - Add improvements:

   • Implement early stopping
   • Add confusion matrix visualization
   • Log model performance metrics

---

Resource Requirements Summary

Resource	Details
GPU	NVIDIA (preferred), min 12GB VRAM
Storage	~50-100GB for expanded dataset
Training Time	6-12 hours for multi-emotion model
Python Libraries	torch, torchvision, ultralytics (already in requirements.txt)
Dataset Sources	FER2013, AffectNet, CK+ databases

---