You can deploy this model and start earning money today!

mattiademartino/
unet-deep-lunar-denoiser

Name: unet-deep-lunar-denoiser
Availability: InStock

Views

No views yet

Lunar Image Denoiser — UNetDeep & Conditional DDPM (Space Data Task 3)

Two models trained to denoise 64×64 lunar surface images from the Space Data Task 3 dataset. Both were selected as the best configurations from an Optuna hyperparameter search (100 trials per architecture/loss combination).

Models

1. UNetDeep — `unet_deep_01` (best model)

A 4-stage encoder–decoder U-Net with skip connections. No attention gates — just clean convolutions at every scale.

Results (150 epochs, 15% validation split)

Metric	Value
Val MSE	0.001758
Val MAE	0.03213
Val PSNR	27.55 dB

Architecture

Layer	In → Out channels	Spatial size (64×64 input)
enc1	1 → 64	64×64
enc2	64 → 128	32×32
enc3	128 → 256	16×16
enc4	256 → 512	8×8
bottleneck	512 → 1024	4×4
dec4	1024 → 512	8×8
dec3	512 → 256	16×16
dec2	256 → 128	32×32
dec1	128 → 64	64×64
out_conv	64 → 1 (sigmoid)	64×64

Dropout: 0.10887 (Dropout2d after each double-conv)
Parameters: ~31 M
Weights file: best_model.pt

Training hyperparameters

Parameter	Value
Loss	MSE
Optimizer	AdamW
Learning rate	9.664e-4
Weight decay	2.472e-3
Batch size	16
Epochs	150
LR schedule	Cosine annealing
Val split	15%
Seed	42

2. Conditional DDPM — `diffusion_01`

A conditional denoising diffusion probabilistic model. The backbone is a 3-stage attention U-Net that takes a 2-channel input [x_t || y] (diffusion state + sensor-noisy conditioning image) and predicts the clean image x0 directly (x0-parameterisation). Timestep t is injected at the bottleneck via a sinusoidal + MLP embedding. Inference uses DDIM deterministic sampling (η=0) for fast generation in as few as 20–50 steps.

Results (150 epochs, 15% validation split, 20 DDIM steps for val)

Metric	Value
Val MSE	0.002654
Val MAE	0.03989
Val PSNR	25.76 dB

Architecture

Component	Details
Input	2-channel: `[x_t \|\| y]`
Encoder stages	3 (features: 64 → 128 → 256)
Bottleneck	512 channels + sinusoidal timestep injection
Decoder	3 stages with soft attention gates on all skips
Output	1 channel, sigmoid → [0, 1]
Diffusion steps	T = 1000 (training), 50 DDIM steps (inference)
Parameters	~14 M

Weights file: best_model_diff.pt

Training hyperparameters

Parameter	Value
Loss	MSE on x0
Optimizer	AdamW
Learning rate	4.248e-4
Weight decay	4.583e-5
Batch size	32
Epochs	150
LR schedule	Cosine annealing
Val split	15%
Seed	42

Dataset

Train: noisy_train_19k_harder.npy / clean_train_19k_harder.npy — ~19k paired 64×64 grayscale patches
Test: noisy_val_1k_harder.npy — 1k noisy patches (blind, no clean labels)
Pixel values normalized to [0, 1]

How to use

UNetDeep (fast, direct inference)

python

import torch
import numpy as np
from huggingface_hub import hf_hub_download

pt_path = hf_hub_download(
    repo_id="mattiademartino/unet-deep-lunar-denoiser",
    filename="best_model.pt",
)

import torch.nn as nn
from typing import Sequence

class DoubleConv(nn.Module):
    def __init__(self, in_ch, out_ch, dropout=0.0):
        super().__init__()
        layers = [
            nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
            nn.Conv2d(out_ch, out_ch, 3, padding=1, bias=False), nn.BatchNorm2d(out_ch), nn.ReLU(inplace=True),
        ]
        if dropout > 0: layers.append(nn.Dropout2d(dropout))
        self.block = nn.Sequential(*layers)
    def forward(self, x): return self.block(x)

class UNetDeep(nn.Module):
    def __init__(self, features=(64,128,256,512), dropout=0.10887):
        super().__init__()
        f = list(features); self.pool = nn.MaxPool2d(2)
        self.enc1, self.enc2 = DoubleConv(1,f[0],dropout), DoubleConv(f[0],f[1],dropout)
        self.enc3, self.enc4 = DoubleConv(f[1],f[2],dropout), DoubleConv(f[2],f[3],dropout)
        self.bottleneck = DoubleConv(f[3],f[3]*2,dropout)
        self.up4, self.dec4 = nn.ConvTranspose2d(f[3]*2,f[3],2,stride=2), DoubleConv(f[3]*2,f[3],dropout)
        self.up3, self.dec3 = nn.ConvTranspose2d(f[3],f[2],2,stride=2),   DoubleConv(f[2]*2,f[2],dropout)
        self.up2, self.dec2 = nn.ConvTranspose2d(f[2],f[1],2,stride=2),   DoubleConv(f[1]*2,f[1],dropout)
        self.up1, self.dec1 = nn.ConvTranspose2d(f[1],f[0],2,stride=2),   DoubleConv(f[0]*2,f[0],dropout)
        self.out_conv = nn.Conv2d(f[0],1,1)
    def forward(self, x):
        e1=self.enc1(x); e2=self.enc2(self.pool(e1)); e3=self.enc3(self.pool(e2)); e4=self.enc4(self.pool(e3))
        b=self.bottleneck(self.pool(e4))
        d4=self.dec4(torch.cat([self.up4(b),e4],1)); d3=self.dec3(torch.cat([self.up3(d4),e3],1))
        d2=self.dec2(torch.cat([self.up2(d3),e2],1)); d1=self.dec1(torch.cat([self.up1(d2),e1],1))
        return torch.sigmoid(self.out_conv(d1))

model = UNetDeep()
model.load_state_dict(torch.load(pt_path, map_location="cpu"))
model.eval()

# single patch inference
noisy = np.load("noisy_val_1k_harder.npy")[0].astype(np.float32) / 255.0
x = torch.from_numpy(noisy).unsqueeze(0).unsqueeze(0)   # (1,1,64,64)
with torch.no_grad():
    denoised = model(x).squeeze().numpy()                # (64,64) in [0,1]

Or use submission.py

bash

# UNet — download weights from Hugging Face and run inference
python submission.py test --hf-model mattiademartino/unet-deep-lunar-denoiser --out predictions.npy

# Diffusion — local weights (auto-searched in project tree)
python submission.py test --model diffusion --weights best_model_diff.pt --out predictions_diffusion.npy

# Diffusion — download from Hugging Face
python submission.py test --model diffusion --hf-model mattiademartino/unet-deep-lunar-denoiser --out predictions_diffusion.npy

# Train from scratch
python submission.py train --data-dir data/ --out-dir results/submission/
python submission.py train --model diffusion --data-dir data/ --out-dir results/submission_diffusion/

mattiademartino/unet-deep-lunar-denoiser