ArjyaDutta-IITM/dlp-nppe3-denoise-sr-x4

DLP NPPE-3 — Low-Light Joint Denoising + 4x Super-Resolution

A from-scratch PyTorch model that takes a noisy, low-light, low-resolution image and produces a denoised, 4x super-resolved image in one forward pass. Built for the IITM BS Data Science DLP course's NPPE-3 competition (metric: PSNR).

Results

The Stage 2 + TTA + calibration configuration is the one shipped in this repo (model.safetensors). The 40-image holdout was never trained on by either stage, so the Stage 1 → Stage 2 gain (+0.17 dB raw) is a genuine improvement, not overfitting to validation data.

Important: the raw model output alone does not reach these scores. 8-way dihedral test-time augmentation and a small affine calibration (alpha, beta in config.json) are both required at inference time — see the code below.

Files

Architecture

LLSRNet (17.1M parameters): a NAFNet-style U-Net denoising trunk (encoder/decoder with channel-attention-free blocks) predicts a clean low-resolution estimate with deep supervision, then an RCAB (residual channel-attention block) stack plus two PixelShuffle stages produce a residual on top of a bicubic upsample of that clean estimate. The bicubic skip means the model starts near a sane baseline (33 dB) at initialization rather than from noise. Full definitions in model_arch.py.

How this was trained

1. 1.Forensics first: the degradation was reverse-engineered directly from paired samples — a 4x4 box-mean downsample, followed by pixel-independent Poisson-Gaussian noise (Var = a·I + b, fit per-image). This let training use the exact clean low-res target as an auxiliary supervision signal and generate unlimited fresh noise realizations for augmentation.
2. 2.Stage 1: trained from scratch on the 1105 training images only, ~180k iterations, progressive patch size (64→80→96), cosine LR schedule, Charbonnier loss transitioning to MSE for the final 20%.
3. 3.Stage 2: fine-tuned from the Stage-1 checkpoint on train + most of the validation set, holding out 40 validation images throughout, so the improvement from folding in validation data could be verified honestly rather than assumed.
4. 4.Inference: 8-way dihedral self-ensembling, then a global affine correction (alpha ≈ 0.998, beta ≈ +0.28) fit on the holdout.

Usage

import json, torch, numpy as np
from PIL import Image
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download
import torch.nn.functional as F

REPO_ID = "ArjyaDutta-IITM/dlp-nppe3-denoise-sr-x4"

# --- download files -------------------------------------------------------
arch_path = hf_hub_download(REPO_ID, "model_arch.py")
cfg_path  = hf_hub_download(REPO_ID, "config.json")
wts_path  = hf_hub_download(REPO_ID, "model.safetensors")

config = json.load(open(cfg_path))
exec(open(arch_path).read(), globals())        # defines LLSRNet etc.

class Cfg:  # minimal stand-in for the training-time CFG object
    pass
cfg = Cfg()
for k, v in config["architecture_args"].items():
    setattr(cfg, k, v)

model = LLSRNet(cfg, config["data_mean"], config["data_std"])
model.load_state_dict(load_file(wts_path))
model.eval().cuda()

# --- 8-way dihedral TTA (required to match reported scores) ---------------
def _aug(x, k):
    if k & 1: x = torch.flip(x, [-1])
    if k & 2: x = torch.flip(x, [-2])
    if k & 4: x = x.transpose(-2, -1)
    return x

def _unaug(x, k):
    if k & 4: x = x.transpose(-2, -1)
    if k & 2: x = torch.flip(x, [-2])
    if k & 1: x = torch.flip(x, [-1])
    return x

def pad_to_multiple(x, m=8):
    h, w = x.shape[-2:]
    ph, pw = (-h) % m, (-w) % m
    if ph or pw:
        x = F.pad(x, (0, pw, 0, ph), mode="reflect")
    return x, h, w

@torch.no_grad()
def predict(lo_img_u8, tta=8):
    # lo_img_u8: (h, w, 3) uint8 low-res, noisy input
    x = torch.from_numpy(lo_img_u8.transpose(2, 0, 1)).float().div_(255.0)
    x = x.unsqueeze(0).cuda()
    scale = cfg.scale
    acc = None
    for k in range(tta):
        xa = _aug(x, k)
        xa, h, w = pad_to_multiple(xa, 8)
        out = model(xa)[0]
        out = _unaug(out[..., :h * scale, :w * scale], k)
        acc = out if acc is None else acc + out
    rgb01 = (acc / tta).clamp(0, 1)[0].permute(1, 2, 0).cpu().numpy()

    # affine calibration on luma (competition metric is luma-based)
    a, b = config["inference"]["calibration_alpha"], config["inference"]["calibration_beta"]
    luma_w = np.array([19595.0, 38470.0, 7471.0]) / 65536.0
    luma = np.clip(rgb01 * 255.0, 0, 255) @ luma_w
    luma_cal = np.clip(np.floor(a * luma + b + 0.5), 0, 255).astype(np.uint8)
    return rgb01, luma_cal   # RGB [0,1] image, and the calibrated luma channel

lo = np.array(Image.open("your_noisy_lowres_image.png").convert("RGB"))
hi_rgb, hi_luma_calibrated = predict(lo, tta=8)

Limitations

• —Trained on a specific synthetic degradation (4x4 box downsample + fitted Poisson-Gaussian noise on very dark images, mean pixel value ≈ 39/255). It is not a general-purpose denoiser or super-resolution model and will likely underperform on images with different noise characteristics, brightness levels, or downsampling kernels.
• —The affine calibration was fit on a 40-image holdout — small enough that its exact values could shift slightly on a different data distribution.
• —Built for a coursework competition; not maintained as a general-purpose tool.