somsubhra04/lowlight-denoise-sr4x

Low-Light Denoising + 4x Super-Resolution (RRDBNet)

Fine-tuned RRDBNet checkpoints for joint denoising + 4x super-resolution of low-light, noisy street-scene images (Cityscapes-derived), trained for the IITM Deep Learning Practice dlp-26t2-nppe3 kaggle contest.

Model Overview

Both checkpoints share the same architecture: the official Real-ESRGAN / BasicSR RRDBNet (23 Residual-in-Residual Dense Blocks, num_feat=64, num_grow_ch=32), initialized from the pretrained RealESRNet_x4plus weights (PSNR-oriented, no GAN loss) and fine-tuned on the competition's low-light noisy → clean high-resolution image pairs.

File	Description	Reported val PSNR
best_model_v4.pth	RRDBNet fine-tuned from RealESRNet_x4plus with a progressive patch-size curriculum, Charbonnier + Sobel gradient loss, EMA weights, final MSE polish phase	39.553 dB
best_model_v4_1.pth	Same as v4, with refined patch schedule / MSE fine-tune fraction	39.571 dB

Both checkpoints store two state dicts:

• model_state — EMA-smoothed weights (use this for inference)
• raw_state — raw (non-EMA) weights at the same training step, kept for reference/resuming only

Other checkpoints from earlier/experimental runs (different architectures, from-scratch runs, etc.) have been moved to other_versions/ for reference.

Training Details

• Base architecture: RRDBNet (num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
• Pretrained init: RealESRNet_x4plus.pth (PSNR-trained, no adversarial loss — chosen over the GAN-trained RealESRGAN_x4plus since this task is scored purely on PSNR)
• Loss: Charbonnier (smooth L1) + a small-weighted Sobel gradient term (0.05), switching to pure MSE for a final polish phase
• Optimizer: AdamW, discriminative learning rates (lower LR for the pretrained body, higher for the newly-adapting upsampling tail), cosine annealing per phase
• EMA: exponential moving average of weights (decay ≈ 0.995–0.999) tracked throughout training; EMA weights consistently outperformed raw weights and are what's stored in model_state
• Patch curriculum: training patches grown across phases (e.g. 80×128 → 112×176 LR patch) rather than a single fixed size, sized to the dataset's native 256×160 LR resolution
• Data: 1,105 training pairs / 267 validation pairs of low-light noisy LR images and their clean 4x-resolution ground truth (see dataset/)

Sample Usage

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
 
# RRDBNet architecture
def make_layer(basic_block, num_basic_block, **kwarg):
    return nn.Sequential(*[basic_block(**kwarg) for _ in range(num_basic_block)])
 
class ResidualDenseBlock(nn.Module):
    def __init__(self, num_feat=64, num_grow_ch=32):
        super().__init__()
        self.conv1 = nn.Conv2d(num_feat, num_grow_ch, 3, 1, 1)
        self.conv2 = nn.Conv2d(num_feat + num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv3 = nn.Conv2d(num_feat + 2 * num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv4 = nn.Conv2d(num_feat + 3 * num_grow_ch, num_grow_ch, 3, 1, 1)
        self.conv5 = nn.Conv2d(num_feat + 4 * num_grow_ch, num_feat, 3, 1, 1)
        self.lrelu = nn.LeakyReLU(0.2, inplace=True)
    def forward(self, x):
        x1 = self.lrelu(self.conv1(x))
        x2 = self.lrelu(self.conv2(torch.cat((x, x1), 1)))
        x3 = self.lrelu(self.conv3(torch.cat((x, x1, x2), 1)))
        x4 = self.lrelu(self.conv4(torch.cat((x, x1, x2, x3), 1)))
        x5 = self.conv5(torch.cat((x, x1, x2, x3, x4), 1))
        return x5 * 0.2 + x
 
class RRDB(nn.Module):
    def __init__(self, num_feat, num_grow_ch=32):
        super().__init__()
        self.rdb1 = ResidualDenseBlock(num_feat, num_grow_ch)
        self.rdb2 = ResidualDenseBlock(num_feat, num_grow_ch)
        self.rdb3 = ResidualDenseBlock(num_feat, num_grow_ch)
    def forward(self, x):
        out = self.rdb1(x); out = self.rdb2(out); out = self.rdb3(out)
        return out * 0.2 + x
 
class RRDBNet(nn.Module):
    def __init__(self, num_in_ch=3, num_out_ch=3, scale=4, num_feat=64, num_block=23, num_grow_ch=32):
        super().__init__()
        self.conv_first = nn.Conv2d(num_in_ch, num_feat, 3, 1, 1)
        self.body = make_layer(RRDB, num_block, num_feat=num_feat, num_grow_ch=num_grow_ch)
        self.conv_body = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)
        self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)
        self.lrelu = nn.LeakyReLU(0.2, inplace=True)
    def forward(self, x):
        feat = self.conv_first(x)
        feat = feat + self.conv_body(self.body(feat))
        feat = self.lrelu(self.conv_up1(F.interpolate(feat, scale_factor=2, mode='nearest')))
        feat = self.lrelu(self.conv_up2(F.interpolate(feat, scale_factor=2, mode='nearest')))
        return self.conv_last(self.lrelu(self.conv_hr(feat)))
 
# load
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ckpt_path = hf_hub_download(repo_id="somsubhra04/lowlight-denoise-sr4x", filename="best_model_v4_1.pth")
model = RRDBNet(num_in_ch=3, num_out_ch=3, scale=4, num_feat=64, num_block=23, num_grow_ch=32).to(device)
ckpt = torch.load(ckpt_path, map_location=device)
model.load_state_dict(ckpt['model_state'])  # EMA weights
model.eval()
 
# inference
with torch.no_grad():
    output = model(lr_tensor.to(device)).clamp(0, 1)  # -> [1, 3, H*4, W*4]

For best results, apply 8-way geometric test-time augmentation (average predictions across the 4 rotations × horizontal flip) — this consistently improved PSNR in evaluation. Round pixel values before casting to uint8 (.round().astype(np.uint8), not plain .astype(np.uint8), which truncates and introduces a small systematic bias).

Repository Structure

best_model_v4.pth 
best_model_v4_1.pth
other_versions/            # earlier / experimental checkpoints (different architectures,
                            # from-scratch runs, etc.) — kept for reference only
dataset/
  train.tar.gz             # 1,105 LR-noisy / HR-clean training pairs
  val.tar.gz                # 267 LR-noisy / HR-clean validation pairs
  test.tar.gz               # 60 LR-noisy test images (no ground truth)

Test-set Results

	PSNR
Highest	39.62946 dB
Ours	39.45618 dB