Skip to content
Pratim/Dasude
08/Computer Vision2 min read

White Balance Regression Model

Outcome

A learned estimator for color temperature and tint, judged on numeric error and on how the corrected image looks.

Computer VisionRegressionPyTorch

EfficientNetV2-S, Computer Vision

View on GitHub
Histograms of absolute error for color temperature and tinterror
Absolute error distributions for color temperature and tint, from the project repository.

Draft write-up

This project trains a CNN to look at an image and regress two numbers: color temperature and tint. Those values drive white balance correction so a warm or cool cast can be neutralised automatically.

Source on GitHub ↗

01Problem

Classic auto white balance relies on heuristics such as gray-world that break on scenes dominated by one color. A learned estimator can use scene context to infer the true illuminant.

02Approach

  1. 01

    Frame it as regression

    Instead of classifying lighting presets, the model predicts continuous temperature and tint, which maps directly onto correction controls.

  2. 02

    Fine-tune EfficientNetV2-S

    A pretrained EfficientNetV2-S backbone gives strong features at a small compute budget. A compact regression head outputs the two targets.

  3. 03

    Augment without breaking labels

    Augmentations are chosen so they do not alter the illuminant the label describes. Color-shifting augmentations are avoided or paired with adjusted labels.

  4. 04

    Evaluate on error that matters

    Beyond MAE per target, results are judged on corrected-image quality, since small numeric errors can be visible or invisible depending on the scene.

03Architecture

A standard training and inference pipeline around a pretrained backbone with a two-output regression head.

01DataIMAGE + (TEMPERATURE, TINT) LABELS02PreprocessRESIZE + LABEL-SAFE AUGMENTATION03BackboneEFFICIENTNETV2-S04HeadREGRESSION LAYERS05Loss & trainingPYTORCH06InferenceCORRECTION MODULE

  1. 01Data

    Image + (temperature, tint) labels

    Dataset loading, normalisation, train/validation split across scenes.

  2. 02Preprocess

    Resize + label-safe augmentation

    Resolution matched to the backbone; geometric augmentations only by default.

  3. 03Backbone

    EfficientNetV2-S

    Pretrained feature extractor fine-tuned end to end.

  4. 04Head

    Regression layers

    Pooled features mapped to two continuous outputs.

  5. 05Loss & training

    PyTorch

    Regression loss on scaled targets, mixed precision, learning-rate schedule.

  6. 06Inference

    Correction module

    Predicted values converted into a white balance adjustment on the image.

  • Continuous temperature and tint prediction
  • Lightweight backbone suitable for image pipelines
  • Per-target error reporting and visual before/after checks

04Challenges

  • Targets sit on very different numeric scales.

    FixStandardised both targets so the loss weights them evenly.

  • Color-based augmentation silently corrupted the labels.

    FixRestricted augmentation to label-preserving transforms.

05Learnings

Outcomes

  • Trained regressor that estimates temperature and tint from a single image
  • Evaluation covering both numeric error and visual correction quality

What I learned

  • Label-aware augmentation matters as much as architecture choice.
  • Numeric error alone is a weak proxy for perceived image quality.

06Next

  • 01Compare against gray-world and other baselines in a table
  • 02Export to ONNX for fast inference