User-2468/mini-unet-colorizer

Mini U-Net Colorizer

A 3,968,892-parameter U-Net that colorizes grayscale photos. Classification-style (Zhang et al., "Colorful Image Colorization"): it predicts a distribution over 236 quantized CIE Lab a/b bins per pixel rather than regressing a single ab value directly, with color-bin loss weights derived from the real training-data distribution (rare/saturated colors weighted higher) so it doesn't just hedge toward desaturated averages. Decode with an annealed mean.

• Status: Final (20 epochs complete)
• Input: L channel normalized as L/50 - 1 -> [-1, 1], shape (1, 256, 256)
• Output: logits over 236 ab bins, shape (236, 256, 256)
• Trained on: johnowhitaker/imagenette2-320 (None), warm-started from User-2468/mini-unet-colorizer
• Loss so far (weighted soft cross-entropy): train 2.4069, val 2.7111

Usage

import numpy as np, torch
from skimage.color import rgb2lab, lab2rgb
from PIL import Image
# paste the SmallUNetColorizer class definition from the training script, then:
model = SmallUNetColorizer.from_pretrained("User-2468/mini-unet-colorizer")
model.eval()

img = Image.open("photo.jpg").convert("RGB").resize((256, 256))
lab = rgb2lab(np.asarray(img).astype("float32") / 255.0)
L = torch.from_numpy(lab[:, :, 0:1] / 50.0 - 1.0).permute(2, 0, 1)[None]

with torch.no_grad():
    logits = model(L)
    ab = model.decode(logits, temperature=0.38)[0].permute(1, 2, 0).numpy()

L_out = (L[0, 0].numpy() + 1) * 50.0
lab_out = np.concatenate([L_out[:, :, None], ab], axis=-1)
rgb_out = np.clip(lab2rgb(lab_out), 0, 1)