Mok777/smart-om-improved-pipeline

Pipeline B+ — Improved SMART-OM Oral Lesion Classification

Model: ConvNeXt-Tiny (IN-22K) + EfficientNet-B2 + GAIN + FiLM + SWA
Task: 4-class oral lesion classification (Normal, VFN, OPMD, OC)
Dataset: SMART-OM, 2469 images
Baseline: Pipeline B (EfficientNet-B2 + MobileViT-XS + GAIN), Macro F1 = 0.7566, VFN F1 = 0.46

---

Summary of Changes (Pipeline B → Pipeline B+)

🏗️ Architecture Changes

Component	Pipeline B	Pipeline B+	Rationale
Primary backbone	EfficientNet-B2 (1408-dim)	ConvNeXt-Tiny IN-22K (768-dim)	ConvNeXt outperforms EfficientNet on medical imaging benchmarks (validated on PCam, BreakHis). IN-22K pretraining provides stronger representations than IN-1K.
Secondary backbone	MobileViT-XS (384-dim)	EfficientNet-B2 (1408-dim)	MobileViT underperforms without strong pretraining on small datasets (−8-10% vs CNNs). EfficientNet-B2 as secondary provides complementary texture features.
Site conditioning	Concatenation (16-dim, 0.88% of fused repr)	FiLM conditioning (32-dim)	Feature-wise Linear Modulation forces backbone features to be conditioned on anatomical site at the architectural level, rather than being marginalized in an MLP.
Fusion normalization	None (raw scale mismatch)	L2 normalize + LayerNorm per branch	Eliminates inter-branch activation scale mismatch. ConvNeXt uses LayerNorm, EfficientNet uses BatchNorm — different output scales.
Fusion head dropout	Default	0.4 (first layer), 0.2 (second layer)	~37M params / ~1700 training samples = 21,800:1 ratio. Higher dropout is critical to prevent overfitting.
Activation function	ReLU (implied)	GELU	Matches ConvNeXt's native activation; smoother gradients for fine-tuning.
Total parameters	~30M	~37M	More parameters but with stronger regularization (dropout, weight decay, SWA).

🎯 GAIN Module Changes

Aspect	Pipeline B	Pipeline B+	Rationale
Supervised classes	{OPMD, OC} only	{VFN, OPMD, OC}	Highest-priority change. VFN has 179 annotations (most per class) but receives zero spatial supervision. VFN F1=46% is the pipeline's defining limitation.
VFN lambda	N/A	0.3	Conservative — lower than OPMD (0.5) since VFN annotation quality is unverified.
Batch coverage	~8.6% of training batches	~19.3% of training batches	More than doubles spatial supervision frequency.
Loss function	MSE (binary mask)	Soft IoU	MSE binarizes CAMs; IoU directly optimizes region overlap.
Mask preprocessing	Binary (hard boundaries)	Gaussian-blurred (σ=5)	Smooth supervision targets match the inherently continuous nature of CAMs.
Degenerate CAM handling	Passes near-zero CAM to MSE	Skips (threshold 1e-8)	Avoids adverse gradients when model hasn't yet learned which features matter.
Hook cleanup	No try/finally	GAIN removed in finally block	Prevents memory leaks on training exceptions.

📊 Loss Function Changes

Aspect	Pipeline B	Pipeline B+	Rationale
Alpha weights	Hand-chosen [0.10, 0.50, 0.90, 1.50]	Effective number of samples (Cui et al., 2019) [0.032, 0.351, 0.501, 3.116]	Principled: α = 1/E_n where E_n = (1-β^n)/(1-β). VFN gets ~11× Normal weight (was 5×). OC gets ~97× Normal weight (was 15×).
Gamma	2.0	2.5	Harder mining on misclassified VFN samples.
Device handling	self.alpha = self.alpha.to(device) in forward()	self.register_buffer('alpha', ...)	Correct semantics — no side effects in forward().
Beta (effective number)	N/A	0.9999	Standard value from Cui et al., 2019.

🏋️ Training Loop Changes

Aspect	Pipeline B	Pipeline B+	Rationale
Optimizer	torch.optim.Adam	torch.optim.AdamW	Decoupled weight decay (Loshchilov & Hutter, 2019). Adam's L2 regularization interacts with adaptive moments — AdamW applies true weight decay.
Weight decay	1e-4	0.05	ConvNeXt recipe; higher WD provides stronger regularization for the high parameter-to-sample ratio.
LR schedule	ReduceLROnPlateau (patience=10)	Cosine annealing with warmup (5 epochs)	Cosine provides more consistent LR decay; warmup prevents early instability. ReduceLROnPlateau with patience=10 and early-stop patience=20 left no room for meaningful LR reduction.
Gradient clipping	None	clip_grad_norm_(max_norm=1.0)	Prevents gradient spikes from focal loss + GAIN combined, especially on high-alpha OC samples.
zero_grad	optimizer.zero_grad()	optimizer.zero_grad(set_to_none=True)	Frees gradient memory rather than zeroing — reduces peak memory.
Training phases	Single phase, differential LR	Phase 1: head-only (5 epochs) → Phase 2: full fine-tune	Decoupled: stabilizes fusion head before backbone adaptation begins.
SWA	None	Epochs 60-80	Stochastic Weight Averaging (Izmailov et al., 2018) averages weights for smoother loss landscape. Typical gain: 0.5-1.5% F1.

📦 Data & Sampling Changes

Aspect	Pipeline B	Pipeline B+	Rationale
Sampling	Uniform random	WeightedRandomSampler (class-balanced)	Each class has equal sampling probability. OC (n=20) appears as often as Normal (n=2145) per epoch.
Augmentation (majority)	Resize→Flip→Rotate→RandomScale→GridDistortion→Resize	RandomResizedCrop→Flip→ShiftScaleRotate→ElasticTransform→Color→CLAHE→Noise→CoarseDropout	Removed double resize and GridDistortion (unrealistic for clinical photos). Added RandomResizedCrop (standard for ImageNet-pretrained models).
Augmentation (minority)	Same as majority	More aggressive: wider crop scale (0.5-1.0), ±45° rotation, ColorJitter, higher dropout	Minority classes need more augmentation diversity to prevent overfitting on few samples.
Mask allocation	Zero mask for all non-annotated (91.4%)	Lazy loading — None for non-annotated	Eliminates waste of allocating 224×224 zero tensors for 91.4% of samples.
drop_last	False	True	Prevents BatchNorm instability from batch-of-2 (1678 mod 16 = 2).
persistent_workers	False	True	Eliminates ~500 worker restarts per fold (100 epochs × 5 folds).
Image validation	None	PIL verify + re-open check	Pre-flight screening catches corrupt images before mid-epoch crashes.
Albumentations API	v1 (var_limit)	v2 (std_range, num_holes_range, etc.)	Correct API for albumentations 2.0+.

🎯 Checkpoint & Ensemble Changes

Aspect	Pipeline B	Pipeline B+	Rationale
Checkpoint metric	Val Macro F1 only	Composite: 0.4×Macro_F1 + 0.6×VFN_F1	Directly optimizes for VFN performance, which is the stated primary goal.
τ-normalization	None	Post-hoc grid search τ∈{0.3, 0.5, 0.7, 0.9, 1.0}	Zero-cost classifier recalibration (Kang et al., ICLR 2020). Corrects majority-class bias in classifier weights. Expected +5-12% tail-class F1.
Ensemble weighting	Equal (1/5 per fold)	Validation composite score weighted	Fold 5 (F1=0.591) contributes less than Fold 2 (F1=0.762).
TTA	None	4-variant: original + hflip + vflip + hflip+vflip	Free accuracy boost (2-4 extra forward passes per test image). Typical +0.5-2% Macro F1.
Variable naming	all_probs stores argmax ints	all_probs stores actual probabilities	Fixes naming bug that caused confusion in threshold tuning notebook.

🔧 Reproducibility Changes

Aspect	Pipeline B	Pipeline B+	Rationale
cudnn.deterministic	Not set	True	Guarantees bit-exact reproducibility.
cudnn.benchmark	Not set (defaults True)	False	Prevents non-deterministic algorithm selection.
PYTHONHASHSEED	Not set	Pinned to seed	Closes last reproducibility gap.
Checkpoint metadata	state_dict only	state_dict + config + fold + epoch + τ + metrics + seed	Enables verification of checkpoint provenance.
Annotation keys	Filename stem (collision risk)	(class_folder, stem) tuple	Eliminates silent collision if two classes share a filename.
Site fallback	Silent → 0 (Dorsal tongue)	Tracked counter + warning	Makes fallback count visible for investigation.

📈 Metrics & Evaluation Changes

Aspect	Pipeline B	Pipeline B+	Rationale
3-class Macro F1	Not reported	Reported alongside 4-class	OC (n=3 test) makes 4-class Macro F1 unreliable. 3-class is the trustworthy metric.
OC metrics	Reported without caveat	Flagged "Only 3 samples — metric unreliable"	Prevents false conclusions from OC F1 = 1.0 or 0.667.
Baseline comparison	Point estimate only	Bootstrap 95% CI (1000 resamples)	Makes improvement claims rigorous for the paper.
GAIN annotated counts	Not printed	Printed per fold	Essential for diagnosing fold-specific GAIN effectiveness.

---

Expected Performance Improvements

Based on the literature (cited above), the expected improvements are:

Metric	Pipeline B	Pipeline B+ (Expected)	Source
Test Macro F1	0.7566	0.80-0.85	ConvNeXt upgrade + class balancing + τ-norm
VFN F1	0.46	0.55-0.65	GAIN extension + effective number weighting + class-balanced sampling + τ-norm
OPMD F1	~0.65	0.70-0.80	Same mechanisms
OC F1	~1.0 (3 samples)	~1.0	Too few samples to reliably measure
Fold variance (std)	0.0593	0.03-0.04	SWA + τ-norm + stronger backbone

Key sources:

• τ-normalization: +5-12% tail-class F1 (Kang et al., ICLR 2020)
• ConvNeXt-T vs EfficientNet-B2: +1-3% on medical imaging (Liu et al., CVPR 2022; validated on PCam, BreakHis)
• Effective number weighting: +2-5% Macro F1 on long-tail datasets (Cui et al., CVPR 2019)
• Class-balanced sampling: +3-8% minority class F1 (standard result)
• SWA: +0.5-1.5% (Izmailov et al., 2018)
• TTA: +0.5-2% at test time (free)

---

Requirements

pip install timm albumentations torch torchvision scikit-learn matplotlib seaborn pandas pillow opencv-python-headless

Tested versions: torch 2.11.0, timm 1.0.26, albumentations 2.0.8

---

Usage

Kaggle

# Update Config paths:
Config.DATA_ROOT = "/kaggle/input/smart-om-dataset"
Config.OUTPUT_DIR = "/kaggle/working/models_improved"

# Run
results = run_pipeline(Config)

Local

Config.DATA_ROOT = "/path/to/SMART-OM"
Config.OUTPUT_DIR = "./models_improved"
Config.FOLD_INDEX_PATH = "./fold_index_assignments.json"
results = run_pipeline(Config)

---

Citation

If you use this pipeline, please cite:

@article{smartom2024,
  title={SMART-OM: A SMARTphone based expert annotated dataset of Oral Mucosa images},
  author={...},
  year={2024}
}

Key methodological references:

• Kang et al., "Decoupling Representation and Classifier for Long-Tailed Recognition", ICLR 2020
• Liu et al., "A ConvNet for the 2020s" (ConvNeXt), CVPR 2022
• Cui et al., "Class-Balanced Loss Based on Effective Number of Samples", CVPR 2019
• Izmailov et al., "Averaging Weights Leads to Wider Optima and Better Generalization", UAI 2018
• Perez et al., "FiLM: Visual Reasoning with a General Conditioning Layer", AAAI 2018