Network API Reference¶

DeepLens provides neural network modules for surrogate modeling of optics and for image reconstruction/restoration.

Surrogate Networks¶

MLP¶

class deeplens.network.MLP(in_features, out_features, hidden_features=64, hidden_layers=3)¶

Multi-layer perceptron producing a normalized output vector (useful for PSF channel normalization).

Parameters:

in_features – Input feature dimension
out_features – Output feature dimension
hidden_features – Hidden layer width
hidden_layers – Number of additional hidden layers

Methods:

forward(x)¶

Forward pass.

Parameters:: x – Input features [B, in_features]
Returns:: Normalized outputs [B, out_features]

MLPConv¶

class deeplens.network.MLPConv(in_features, ks, channels=3, activation='relu')¶

MLP encoder + convolutional decoder that maps a condition vector to a spatial kernel/map. Useful for high-frequency PSF prediction.

Parameters:

in_features – Input feature dimension (e.g., condition)
ks – Output spatial size (kernel size)
channels – Number of output channels
activation – Activation in bottleneck (“relu” or “sigmoid”)

Methods:

forward(x)¶

Generate spatial output.

Parameters:: x – Condition vector [B, in_features]
Returns:: Spatial map [B, channels, ks, ks]

Siren (sine linear layer)¶

class deeplens.network.Siren(dim_in, dim_out, w0=1.0, c=6.0, is_first=False, use_bias=True, activation=None)¶

Sine-activated linear layer used to build SIREN-style implicit networks.

Parameters:

dim_in – Input dimension
dim_out – Output dimension
w0 – Sine frequency multiplier
c – Initialization constant
is_first – Whether this is the first SIREN layer
use_bias – Use bias term
activation – Optional activation module (defaults to sine)

Methods:

forward(x)¶

Parameters:: x – Inputs [B, dim_in]
Returns:: Outputs [B, dim_out]

ModulateSiren¶

class deeplens.network.ModulateSiren(dim_in, dim_hidden, dim_out, dim_latent, num_layers, image_width, image_height, w0=1.0, w0_initial=30.0, use_bias=True, final_activation=None, outermost_linear=True)¶

SIREN-based synthesizer modulated by a latent vector (FiLM-like). Internally samples a fixed coordinate grid in [-1, 1] for the given spatial size and produces an image/map.

Parameters:

dim_in – Coordinate input dimension (typically 2)
dim_hidden – Hidden width in synthesizer SIREN
dim_out – Number of output channels
dim_latent – Latent vector dimension used by the modulator
num_layers – Number of SIREN layers in synthesizer/modulator
image_width – Output width
image_height – Output height
w0 – Frequency scale for hidden SIREN layers
w0_initial – Frequency scale for the first SIREN layer
use_bias – Use bias terms
final_activation – Optional final activation
outermost_linear – If True, ends with Linear; otherwise another SIREN layer

Methods:

forward(latent)¶

Parameters:: latent – Latent vector [B, dim_latent]
Returns:: Outputs [B, dim_out, image_height, image_width]

PSFNet_MLPConv¶

class deeplens.network.surrogate.PSFNet_MLPConv(in_chan=2, kernel_size=128, out_chan=3, latent_dim=4096, latent_channels=16)¶

Combined MLP conditioner and convolutional decoder for spatially varying PSF modeling.

Parameters:

in_chan – Conditioner input dimension (e.g., 2 for (r, z))
kernel_size – Output PSF size (assumes powers of 2 per implementation)
out_chan – Output channels
latent_dim – Conditioner output (flatten) dimension
latent_channels – Channels used when reshaping latent to a feature map

Methods:

forward(x)¶

Parameters:: x – Conditioner inputs [B, in_chan]
Returns:: PSF tensor [B, out_chan, kernel_size, kernel_size]

Reconstruction Networks¶

UNet¶

class deeplens.network.UNet(in_channels=3, out_channels=3)¶

Lightweight UNet variant for image restoration.

Parameters:

in_channels – Input channels
out_channels – Output channels

Methods:

forward(x)¶

Parameters:: x – Input image [B, C, H, W]
Returns:: Restored image [B, C, H, W]

NAFNet¶

class deeplens.network.NAFNet(in_chan=3, out_chan=3, width=32, middle_blk_num=1, enc_blk_nums=[1, 1, 1, 28], dec_blk_nums=[1, 1, 1, 1])¶

Nonlinear Activation Free Network for image restoration.

Parameters:

in_chan – Input channels
out_chan – Output channels
width – Base width
middle_blk_num – Number of middle blocks
enc_blk_nums – Encoder blocks per scale
dec_blk_nums – Decoder blocks per scale

Methods:

forward(inp)¶

Parameters:: inp – Input image [B, C, H, W]
Returns:: Restored image [B, C, H, W]

Restormer¶

class deeplens.network.Restormer(inp_channels=3, out_channels=3, dim=48, num_blocks=[4, 6, 6, 8], num_refinement_blocks=4, heads=[1, 2, 4, 8], ffn_expansion_factor=2.66, bias=False, LayerNorm_type='WithBias', dual_pixel_task=False)¶

Transformer-based restoration model.

Parameters:

inp_channels – Input channels
out_channels – Output channels
dim – Base dimension
num_blocks – Blocks per encoder/decoder stage
num_refinement_blocks – Refinement blocks at the end
heads – Attention heads per scale
ffn_expansion_factor – FFN expansion factor
bias – Whether to use bias
LayerNorm_type – ‘WithBias’ or ‘BiasFree’
dual_pixel_task – Enable dual-pixel defocus deblurring setting

Methods:

forward(inp_img)¶

Parameters:: inp_img – Input image [B, C, H, W]
Returns:: Restored image [B, C, H, W]

SwinIR¶

class deeplens.network.reconstruction.SwinIR(img_size=64, patch_size=1, in_chans=3, embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6], window_size=7, mlp_ratio=4.0, qkv_bias=True, qk_scale=None, drop_rate=0.0, attn_drop_rate=0.0, drop_path_rate=0.1, norm_layer=nn.LayerNorm, ape=False, patch_norm=True, use_checkpoint=False, upscale=2, img_range=1.0, upsampler='', resi_connection='1conv')¶

Swin Transformer-based restoration/upsampling model.

Parameters:

img_size – Input image size (int or tuple)
patch_size – Patch size
in_chans – Input channels
embed_dim – Embedding dimension
depths – Depth per stage
num_heads – Attention heads per stage
window_size – Window size
mlp_ratio – MLP ratio
upscale – Upscale factor (2/3/4/8; 1 for denoising/JPEG)

Loss Functions¶

PSNRLoss¶

class deeplens.network.PSNRLoss(loss_weight=1.0, reduction='mean', toY=False)¶

PSNR-based loss; larger PSNR corresponds to lower loss.

Methods:

forward(pred, target)¶

Parameters:

pred – Predicted images [B, C, H, W]
target – Ground truth images [B, C, H, W]

Returns:

Loss scalar

SSIMLoss¶

class deeplens.network.SSIMLoss(window_size=11, size_average=True)¶

SSIM-based loss returning 1 - SSIM.

Methods:

forward(pred, target)¶

Parameters:

pred – Predicted images [B, C, H, W]
target – Ground truth images [B, C, H, W]

Returns:

Loss scalar (1 - SSIM)

PerceptualLoss¶

class deeplens.network.PerceptualLoss(device=None, weights=[1.0, 1.0, 1.0, 1.0, 1.0])¶

VGG16 feature-based perceptual loss using layers relu1_2, relu2_2, relu3_3, relu4_3, relu5_3.

Methods:

forward(x, y)¶

Parameters:

x – Predicted images [B, 3, H, W]
y – Target images [B, 3, H, W]

Returns:

Perceptual loss scalar

Datasets¶

ImageDataset¶

class deeplens.network.ImageDataset(img_dir, img_res=None)¶

Basic image dataset loader with augmentation and ImageNet-style normalization.

Methods:

__getitem__(idx)¶

Returns:: Tensor image [C, H, W]

__len__()¶

Returns:: Number of images

PhotographicDataset¶

class deeplens.network.PhotographicDataset(img_dir, img_res=(512, 512), iso_range=(100, 400), is_train=True)¶

Image dataset that also samples ISO and field center for simulation-driven training.

Methods:

__getitem__(idx)¶

Returns:: Dict with keys img (Tensor [C, H, W]), iso (float), iso_scale (int), field_center (Tensor [2])

__len__()¶

Returns:: Number of images

Dataset download helpers¶

deeplens.network.dataset.download_bsd300(destination_folder='./datasets')¶: Download BSDS300 to destination_folder and return the image directory path.

deeplens.network.dataset.download_div2k(destination_folder)¶: Download and extract DIV2K (HR) splits into destination_folder.

deeplens.network.dataset.download_flick2k(destination_folder='./datasets')¶: Download and extract Flickr2K into destination_folder (via Hugging Face).

deeplens.network.dataset.download_div8k(destination_folder='./datasets')¶: Download and extract DIV8K into destination_folder (via Hugging Face).

Examples¶

PSF modeling with PSFNet_MLPConv¶

import torch
from deeplens.network.surrogate import PSFNet_MLPConv

# (r, z) conditioner -> PSF
model = PSFNet_MLPConv(in_chan=2, kernel_size=128, out_chan=3)
rz = torch.tensor([[0.5, -5000.0], [-0.3, -2000.0]])  # [B=2, 2]
psf = model(rz)  # [2, 3, 128, 128]

Image Restoration with NAFNet¶

import torch
from deeplens.network import NAFNet

model = NAFNet(in_chan=3, out_chan=3)
degraded = torch.rand(1, 3, 256, 256)
restored = model(degraded)

Combined Loss¶

import torch
from deeplens.network import PSNRLoss, SSIMLoss, PerceptualLoss

class CombinedLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.psnr = PSNRLoss()
        self.ssim = SSIMLoss()
        self.perc = PerceptualLoss()
    def forward(self, pred, target):
        loss = 0.0
        # Lower MSE -> higher PSNR -> lower PSNRLoss
        loss += 1.0 * self.psnr(pred, target)
        # SSIMLoss already returns (1 - SSIM)
        loss += 0.5 * self.ssim(pred, target)
        loss += 0.1 * self.perc(pred, target)
        return loss

DeepLens

Navigation

Related Topics

Network API Reference¶

Surrogate Networks¶

MLP¶

MLPConv¶

Siren (sine linear layer)¶

ModulateSiren¶

PSFNet_MLPConv¶

Reconstruction Networks¶

UNet¶

NAFNet¶

Restormer¶

SwinIR¶

Loss Functions¶

PSNRLoss¶

SSIMLoss¶

PerceptualLoss¶

Datasets¶

ImageDataset¶

PhotographicDataset¶

Dataset download helpers¶

Examples¶

PSF modeling with PSFNet_MLPConv¶

Image Restoration with NAFNet¶

Combined Loss¶

See Also¶