pvjosue / WaveBlocks

Optimization framework developed in Pytorch, allowing calibration, and joint optimization of optics and deep learning models.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Apache License Google Scholar PDF

WaveBlocks

About

A polimorfic optimization framework for Fluorescence Microscopy based on Pytorch, capable of:

  • Simulation.
  • Calibration.
  • PSF engineering.
  • Joint optimization of optics and reconstruction, segmentation and arbitrary optimization tasks.
  • Plug and play to Pytorch neural networks.

Wave-optics involves complex diffraction integrals that when stacking many optical elements become very hard to derivate and optimize. Here we take advantage of the linearity of the possible operations and modularize them into blocks, enabling building arbitrarily large systems.

By building a WaveBlocks microscope similar to the one in your lab you can first calibrate the unknown parameters, like an accurate PSF, distance between elements, and even errors caused by aberrations and SLM diffraction. Once with the capability of simulating images similar to your microscope you can plug in any Pytorch based network to the imaging system and find and optimize the optimal optical parameters together with your network.

Workflow

  • Each optical element in a microscope (lenses, propagation through a medium, cameras, PSFs, Spatial Light modulator, apertures, Micro-lens arrays, etc) is represented by a block (Lego like) that can be assembled in any order, each block has as an input/output a complex Wave-Front and processes the input given the wave-optics behavior of the block.

  • Most of the variables of the blocks are optimizable, allowing straight forward tasks like calibrating distances between optical elements, correction factors, compute real PSFs, all based on real data.

  • All the blocks are based on a OpticBlock class. Which takes care of gathering the parameters to optimize selected by the user and return them to be fed to the optimizer (Adam, SGD, etc).

  • Currently the main objective + tube-lens PSF is imported from an external file. Use Matlab, Fiji or your chosen software to generate the PSF. Its shape should be [1,nDepths,dim1,dim2,2] (the last dimension gathers the real and imaginary parts.

Current implemented Waveblocks

OpticConfig (nn.Module)

Class containing global information about the optical setup, such as wavelength, wave-number and refractive index. It is also used for populating parameters across the microscope.

OpticalBlock (nn.Module)

Base Class of the framework, containing an OpticConfig and members_to_learn, identifying which Parameters from the current block should be optimized. These are provided by the user.

WavePropagation (OpticalBlock)

Uses the Rayleight-Sommerfield integral to propagate an incoming complex wave-front a given distance.

Lens (OpticalBlock)

This class can propagate a wave from the focal plane to the back focal plane or from any point i in front of the lens, to any point o in the back of the lens.

DiffractiveElement (OpticalBlock)

Class simulating any optical element that modifies either the amplitude or phase of an incoming wave-front, such as: apertures, masks and phase-masks.

MicroLensArray (OpticalBlock)

This class contains the functionality for space invariant PSF's like in micro lens arrays, used in Light-field microscopes.

Camera (OpticalBlock)

Computes the intensity of the field hitting the camera and convolves it with an observed object. Works with space variant and invariant PSFs.

Use Waveblocks locally

Generating distribution archives

To distribute WaveBlocks on PyPa it has to be packaged this can be done with the following commands. For a more detailed explanation please follow this guide.

python3 -m pip install --upgrade build
python3 -m build
pip install -r requirements.txt

Install WaveBlocks

Since WaveBlocks is currently not published on PyPa but its still packaged with pip one has to install it locally to use it. The following command has to be executed in the /WaveBlocks folder abd will install the package locally in an editable mode.

pip3 install -e ["PathToWaveBlocksFolder"]

Microscope example

Microscope Class derived from Optical block, simulates the behavior of a microscope

import waveblocks as ob
import torch
import torch.nn as nn
    
class Microscope(ob.OpticBlock):
    def __init__(self, psfIn, vars_to_learn, optical_config):
        super(Microscope, self).__init__(optical_config, vars_to_learn)

        # create parameter for input PSF
        self.psf = nn.Parameter(psfIn, requires_grad=True)
        
        # Fetch pixel size and sampling rate in general
        self.sampling_rate = optical_config.sensor_pitch
        
        # Lateral size of PSF
        field_length = psfIn.shape[2]

        # Create Wave-Propagation block, to defocus the PSF to a given depth
        self.wave_prop = ob.WavePropagation(self.optic_config, [], self.sampling_rate, optical_config.minDefocus, field_length)
        
        # Create Camera Block
        self.camera = ob.Camera(self.optic_config, [], self.sampling_rate)

    def forward(self, realObject):
        # Fetch PSF
        psf = self.psf

        # Defocus PSF given wave_prop.propagation_distance
        defocused_psf = self.wave_prop(psf)
        
        # Compute PSF irradiance and convolve with object
        finalImg,psf = self.camera(realObject, defocused_psf)
        return finalImg

Asumptions

The following design assumptions and functionalites are taken into account:

  • All the blocks work with a wave-optics approach of light simulation, appropiate for difracted limited systems like microscopes.
  • Light sources are assumed to be incoherent (as in fluorescence microscopy) (coherent is not yet iplemented but will be in the future for lasers and holography).
  • The scalar diffraction theory is assumed (paraxial approximation, small angles and low NA)

Citing this work

@article{pageWaveBlocks2020,
    author = {Page, Josue and Favaro, Paolo},
    booktitle = {IEEE International Conference on Image Processing (ICIP)},
    title = {{L}earning to {M}odel and {C}alibrate {O}ptics via a {D}ifferentiable {W}ave {O}ptics {S}imulator},
    year = {2020}
}

Contact

Josue Page - josue.page@tum.de

About

Optimization framework developed in Pytorch, allowing calibration, and joint optimization of optics and deep learning models.


Languages

Language:Jupyter Notebook 66.1%Language:Python 33.9%