core/model.py

"""Small DnCNN-style residual denoiser (predicts the noise; clean = noisy - noise)."""
from __future__ import annotations

import numpy as np


def build_net(depth: int = 6, ch: int = 32):
    import torch.nn as nn
    layers = [nn.Conv2d(1, ch, 3, padding=1), nn.ReLU(inplace=True)]
    for _ in range(depth - 2):
        layers += [nn.Conv2d(ch, ch, 3, padding=1), nn.BatchNorm2d(ch), nn.ReLU(inplace=True)]
    layers += [nn.Conv2d(ch, 1, 3, padding=1)]
    return nn.Sequential(*layers)


def load_model(path: str):
    import torch
    net = build_net()
    net.load_state_dict(torch.load(path, map_location="cpu"))
    net.eval()
    return net


def denoise(net, noisy: np.ndarray) -> np.ndarray:
    import torch
    with torch.no_grad():
        x = torch.from_numpy(noisy[None, None].astype(np.float32))
        residual = net(x)[0, 0].cpu().numpy()
    return np.clip(noisy - residual, 0, 1).astype(np.float32)