U-Net model for low-light to daylight image enhancement. Introduces two targeted fixes for the gray/washed-out color problem that affected v1 and v2: residual learning (model predicts a delta rather than the full image) and a differentiable YCbCr color loss that explicitly penalises desaturated outputs.
Why v3 Exists
v1 and v2 produced structurally correct but washed-out, gray images. Two compounding causes:
Sigmoid → 0.5 on uncertainty. When the model is unsure about a pixel's correct color, Sigmoid activations converge to 0.5 — mid-gray. A blue sky the model hasn't confidently learned comes out gray.
Color-blind losses. L1 and MS-SSIM both accept desaturated outputs that match luminance/structure but miss chrominance. A gray sky can have good SSIM with a blue sky if luminance matches.
This means the model only needs to learn what changes between night and day. Background structure is preserved for free. Tanh activations don't saturate to a mid-gray mean — they saturate to ±1, so uncertain predictions push toward the input rather than toward gray.
Loss Function: ColorLoss
ColorLoss converts predictions to YCbCr (BT.601 coefficients, pure PyTorch — no new dependencies) and penalises chrominance errors 2× relative to luminance:
Channel-wise MSE (float64, range [0, 1]) on validation set and held-out eval pair:
Metric
v1 (best.pt)
v3 (best_v3.pt)
Val MSE avg
0.028953
0.050844
Val MSE — R
0.028700
0.050646
Val MSE — G
0.025959
0.046377
Val MSE — B
0.032200
0.055508
Eval MSE avg (night.jpg)
0.038925
0.075302
Note on higher MSE: v3's val and eval MSE are higher than v1. This is expected and not a regression. The loss function explicitly trades raw pixel accuracy for chrominance correctness — a gray sky that matches luminance structure will score better than a blue sky on pure MSE, but worse on ColorLoss. v3 is optimised for color fidelity; MSE is reported for comparability only. Visual inspection is the meaningful test.
Usage
python
1import torch
2from PIL import Image
3from torchvision import transforms
4from huggingface_hub import hf_hub_download
5from model import UNet # download model.py from this repo67# Load model — note residual=True for v38ckpt_path = hf_hub_download("tyakovenko/night-to-day-enhancement-model-v3","best_v3.pt")9ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)10base_filters = ckpt.get("args",{}).get("base_filters",16)11model = UNet(base_filters=base_filters, residual=True)12model.load_state_dict(ckpt["model"])13model.eval()1415# Inference16img = Image.open("night.jpg").convert("RGB")17x = transforms.ToTensor()(img).unsqueeze(0)# [1, 3, H, W]1819# Pad to multiple of 16 (required by 4-level U-Net)20defpad_to_multiple(t, m=16):21 _, _, h, w = t.shape
22 ph =(m - h % m)% m
23 pw =(m - w % m)% m
24return torch.nn.functional.pad(t,(0, pw,0, ph), mode="reflect"), h, w
2526x_padded, orig_h, orig_w = pad_to_multiple(x)27with torch.no_grad():28 out = model(x_padded)29out = out[:,:,:orig_h,:orig_w]3031result = transforms.ToPILImage()(out.squeeze(0).clamp(0,1))32result.save("enhanced.jpg")
Model Lineage
Version
Key change
Val MSE
v1
Baseline — U-Net, MSE loss, Transient Attributes
0.028953
v1-extended
+ LOL fine-tuning
0.027752
v2
L1 + MS-SSIM loss
—
v3
Residual U-Net + YCbCr ColorLoss
0.050844
See the project repo for full training scripts and architecture details.
Limitations
v3 trades pixel-level MSE for chrominance accuracy — raw MSE is higher than v1
Residual learning helps but 20 epochs from a Sigmoid warm-start may not be enough for full recalibration; more training likely beneficial
Staged training (warm-start longer before introducing ColorLoss) is a promising next step
