Robust Mask Wearing Detection via Label Correction, Class Rebalancing, and Coordinate Attention-Enhanced YOLOv5s

A joint data-model co-optimization strategy for three-state mask wearing detection. On the data side, we correct faulty annotations and apply class-balancing augmentation. On the model side, we integrate Coordinate Attention (CA) post-SPPF in YOLOv5s with two-stage transfer learning and SGD optimization. The final model achieves mAP@0.5 = 0.910, mAP@0.5:0.95 = 0.646, and Recall = 0.878, substantially outperforming SSD+LeNet-5, vanilla YOLOv5s, and other attention variants (SE, CBAM).

---

Dataset

The original Kaggle Face Mask Detection dataset (853 images) exhibits severe class imbalance and annotation noise. We apply a two-step refinement:

Label Correction: Manual inspection fixes mislabeled classes, over-sized/offset bounding boxes, and missed faces in crowded scenes.

Class Rebalancing: Supplement minority-class samples from external data, then apply Copy-Paste augmentation. Cropped foreground patches are randomly pasted onto backgrounds with automatically generated bounding boxes.

Category	Before	After
with_mask	2,151	2,986
without_mask	633	2,519
mask_worn_incorrect	83	806
Total	2,867	6,311

Split: 70/15/15, stratified by dominant label per image, seed=42.

Classes

ID	Name	Description
0	with_mask	Mask correctly covering nose and mouth
1	without_mask	No mask worn
2	mask_worn_incorrect	Mask worn incorrectly (not covering nose/mouth)

Metrics

Metric	Definition
mAP@0.5	Mean Average Precision at IoU >= 0.5
mAP@0.5:0.95	Averaged over IoU thresholds 0.5 to 0.95 (step 0.05)
Precision	TP / (TP + FP)
Recall	TP / (TP + FN)

---

Method

Architecture

We insert a Coordinate Attention (CA) module after the SPPF layer of YOLOv5s:

YOLOv5s:        ... -> SPPF(layer 9)  -> Head(layer 10)
YOLOv5s-CA:     ... -> SPPF(layer 9)  -> CoordAtt(layer 10) -> Head(layer 11)

CA decomposes global pooling into 1D directional encodings (H-pool and W-pool), preserving precise positional information. This coordinate-aware design distinguishes valid facial regions from padded backgrounds -- critical for disambiguating densely packed faces.

Training Strategy

Naively inserting CA creates an optimization barrier: the randomly initialized attention layer disrupts gradient flow into the backbone. We solve this with two-stage transfer learning:

Stage	Description	Epochs	LR
Stage 1	Train full YOLOv5s from scratch (SGD)	~200	lr0=0.01
Stage 2	Load backbone weights; randomly init CoordAtt; fine-tune all	300	lr0=0.01

SGD consistently outperforms AdamW across all configurations. Fixed hyperparameters: box=7.5, cls=1.0, dfl=1.5, conf=0.25, batch=16, imgsz=640.

---

Results

All results on test set. Training: NVIDIA V100 (16 GB) via SLURM.

Overall Performance

Model	mAP@0.5	mAP@0.5:0.95	Precision	Recall
B1 (SSD + LeNet-5)	0.159	0.053	0.460	0.196
YOLOv5s	0.861	0.588	0.869	0.803
YOLOv5s-CA (Ours)	0.910	0.646	0.906	0.878

Per-class AP@0.5: B1 nearly fails on mask_worn_incorrect (0.001). YOLOv5s-CA delivers the largest gains on minority classes -- mask_worn_incorrect surges from 0.807 to 0.930 (+12.3%), Recall from 0.731 to 0.870 (+13.9%).

Ablation Study 1: Optimizer Selection

Model	Optimizer	mAP@0.5	mAP@0.5:0.95	Precision	Recall
YOLOv5s-CA	AdamW	0.852	0.571	0.920	0.724
YOLOv5s-CA	SGD	0.910	0.646	0.906	0.878

Ablation Study 2: Transfer Learning

Model	Init	mAP@0.5	mAP@0.5:0.95	Precision	Recall
YOLOv5s-CA	Scratch	0.904	0.643	0.916	0.826
YOLOv5s-CA	Transfer	0.910	0.646	0.906	0.878

CA from scratch underperforms vanilla YOLOv5s (0.904 vs 0.909). Transfer learning is required to unlock CA's potential.

Ablation Study 3: CA Insertion Position

Model	Position	mAP@0.5	mAP@0.5:0.95	Precision	Recall
YOLOv5s-CA	Mid-C3	0.880	0.621	0.850	0.845
YOLOv5s-CA	Post-SPPF	0.910	0.646	0.906	0.878

Post-SPPF provides +3.0 pp mAP@0.5. SPPF-aggregated multi-scale features offer the richest context for attention recalibration.

Ablation Study 4: Attention Mechanism Comparison

Model	mAP@0.5	mAP@0.5:0.95	Precision	Recall
YOLOv5s (vanilla)	0.861	0.588	0.869	0.803
YOLOv5s + SE	0.884	0.627	0.915	0.850
YOLOv5s + CBAM	0.907	0.642	0.907	0.836
YOLOv5s + CA	0.910	0.646	0.906	0.878

Only CA's coordinate-aware decomposition provides precise positional encoding. SE lacks spatial guidance; CBAM's global pooling loses coordinates.

Dataset Refinement Impact

Metric	Original (853 img)	Refined (2060 img)
mAP@0.5	0.856	0.861
mAP@0.5:0.95	0.589	0.589
Precision	0.902	0.870
Recall	0.786	0.803

Label correction and class rebalancing improve minority-class Recall by +6.6 pp on without_mask.

---

Project Structure

CS182_project_mask_detection/
├── baselines/              # B1 two-stage baseline
│   ├── b1_hog_svm_detector.py  # LeNet-5 classifier + two-stage evaluation
│   └── two_stage_common.py     # Face proposals, crop, eval utilities
├── detectors/              # Unified training entry point
│   └── train.py                # --model {yolov5s|yolov8s|yolo11s|yolov5s-ca|yolov5s-cbam|yolov5s-se|yolov5s-ca-paper|yolov5s-twostage}
├── configs/                # Model architecture YAMLs
│   ├── mask_data.yaml
│   ├── yolov5s-ca.yaml         # CA post-SPPF (best)
│   ├── yolov5s-ca-mid.yaml     # CA mid-backbone (ablation)
│   ├── yolov5s-ca-paper.yaml   # CA paper architecture
│   ├── yolov5s-cbam.yaml
│   └── yolov5s-se.yaml
├── fmd/                    # Shared library
│   ├── attention.py            # CoordAtt, CBAM, SE modules
│   ├── constants.py            # Paths, class names
│   ├── det_eval.py             # mAP computation, IoU matching
│   ├── device.py               # GPU/CPU selection
│   ├── io_utils.py             # JSON I/O
│   ├── metrics.py              # COCO-style metrics
│   └── voc.py                  # Pascal VOC XML parser
├── scripts/                # Evaluation & visualization
│   ├── eval.py                 # Unified eval: --compare / --ablate / eval
│   ├── sweep.py                # Hyperparameter sweep: conf / loss
│   ├── plot_pr_curve.py        # PR curve analysis
│   ├── heatmaps.py             # Attention heatmap comparison (4 models)
│   ├── generate_visualizations.py  # Report figures
│   ├── visualize_three_models.py   # B1/M2/M3 comparison
│   ├── run_specialized_experiments.py  # Dense crowd & low-light
│   ├── test_dense_4models.py   # Dense 4-model benchmark
│   ├── ca_failure_analysis.py  # Failure mode analysis
│   ├── download_dataset.py     # Kaggle download
│   ├── prepare_original_dataset.py  # Original (non-augmented) dataset prep
│   ├── split_dataset.py        # Stratified 70/15/15 split
│   └── voc_to_yolo.py          # VOC XML -> YOLO txt
├── slurm/                  # SLURM job scripts (exp_* = paper experiments)
│   ├── train.slurm              # Unified training entry
│   ├── exp_dataset_refinement.slurm  # Dataset refinement comparison
│   ├── exp_pr_curve.slurm       # Appendix A: PR curve sweep
│   ├── exp_attention_heatmaps.slurm  # Figure: attention heatmap comparison
│   ├── exp_dense_lowlight.slurm # Dense crowd & low-light experiments
│   ├── exp_dense_4model.slurm   # Dense 4-model comparison
│   ├── exp_viz_3models.slurm    # B1/M2/M3 visualization
│   └── exp_failure_analysis.slurm  # Appendix D: failure mode analysis
├── requirements.txt
└── README.md

---

Quick Start

Environment

conda create -n fmd python=3.10 -y
conda activate fmd
pip install ultralytics torch pandas matplotlib seaborn opencv-python pyyaml

Data Preparation

python scripts/download_dataset.py       # Requires Kaggle API token
python scripts/split_dataset.py          # 70/15/15 stratified split
python scripts/voc_to_yolo.py            # VOC XML -> YOLO txt

Training

All models via single entry point detectors/train.py:

python detectors/train.py --model <MODEL> [--epochs 300] [--pretrained auto] ...

--model	Description
yolov5s	Vanilla YOLOv5s baseline
yolov8s	YOLOv8s
yolo11s	YOLO11s
yolov5s-ca	YOLOv5s + CA post-SPPF (best)
yolov5s-cbam	YOLOv5s + CBAM
yolov5s-se	YOLOv5s + SE
yolov5s-twostage	Two-stage fine-tuned YOLOv5s

Reproduce Key Experiments

# Ablation 1: Optimizer (SGD vs AdamW)
python detectors/train.py --model yolov5s-ca --epochs 300 --patience 50 --name ca_sppf_sgd
python detectors/train.py --model yolov5s-ca --epochs 300 --patience 50 --optimizer AdamW --name ca_sppf_adamw

# Ablation 2: Transfer learning (best config)
python detectors/train.py --model yolov5s --epochs 200 --name yolov5s_mask              # Stage 1
python detectors/train.py --model yolov5s-ca --pretrained auto --epochs 300 --patience 50 --name ca_sppf_pretrain_sgd  # Stage 2

# Ablation 3: CA position
python detectors/train.py --model yolov5s-ca --cfg configs/yolov5s-ca-mid.yaml --pretrained auto --epochs 300 --patience 50

# Ablation 4: Attention variants
python detectors/train.py --model yolov5s-cbam --pretrained auto --epochs 300 --patience 50
python detectors/train.py --model yolov5s-se --pretrained auto --epochs 300 --patience 50

# Dataset refinement comparison
sbatch slurm/exp_dataset_refinement.slurm yolov5s    # Original dataset

# SLURM cluster
sbatch slurm/train.slurm yolov5s-ca --pretrained auto --epochs 300 --patience 50

Evaluation

# Multi-model comparison
python scripts/eval.py --compare

# Ablation summary with dimensional deltas
python scripts/eval.py --ablate

# Single model
python scripts/eval.py eval --weights results/detectors/ca_sppf_pretrain_sgd/weights/best.pt --split test

Visualization & Analysis

# PR curve
python scripts/plot_pr_curve.py --split test --iou 0.5
# or: sbatch slurm/exp_pr_curve.slurm

# Attention heatmaps
python scripts/heatmaps.py --device cuda --max-samples 6
# or: sbatch slurm/exp_attention_heatmaps.slurm

# Failure analysis
python scripts/ca_failure_analysis.py --device cuda --conf 0.25 --max-cases 6
# or: sbatch slurm/exp_failure_analysis.slurm

# Dense crowd & low-light
python scripts/run_specialized_experiments.py --device cuda --splits test --dense-min-faces 8
# or: sbatch slurm/exp_dense_lowlight.slurm

# Dense 4-model comparison
python scripts/test_dense_4models.py --device cuda --splits train,val,test --dense-min-faces 5
# or: sbatch slurm/exp_dense_4model.slurm

---

References

1. Hou, Q., Zhou, D., & Feng, J. (2021). Coordinate Attention for Efficient Mobile Network Design. CVPR, 13713--13722.
2. Hu, J., Shen, L., & Sun, G. (2018). Squeeze-and-Excitation Networks. CVPR, 7132--7141.
3. Woo, S., Park, J., Lee, J.-Y., & Kweon, I. S. (2018). CBAM: Convolutional Block Attention Module. ECCV, 3--19.
4. Jocher, G. et al. (2022). ultralytics/yolov5: v7.0. Zenodo. doi:10.5281/zenodo.7347926
5. Ultralytics YOLOv8 / YOLO11. https://github.com/ultralytics/ultralytics
6. Kaggle Face Mask Dataset. https://www.kaggle.com/datasets/andrewmvd/face-mask-detection
7. Properly Wearing Masked Detect Dataset. https://github.com/ethancvaa/Properly-Wearing-Masked-Detect-Dataset