ZRNet: Physics-Informed Network with Frequency-Aware Learning for Optical Aberration Correction

Installation

This implementation is based on BasicSR which is an open-source toolbox for image/video restoration tasks, NAFNet, Restormer and Multi Output Deblur

# CUDA 11.1
pip install torch==1.10.1+cu111 torchvision==0.11.2+cu111 torchaudio==0.10.1 -f https://download.pytorch.org/whl/cu111/torch_stable.html

pip install -r requirements.txt
python setup.py develop --no_cuda_ext

Data Preparation

Download CytoImageNet dataset from Kaggle.

Add full_path column to the metadata.csv file.

import pandas as pd
import os

input_filename = 'metadata.csv'
# The base directory to prepend to the path
base_dir = '/db'

df = pd.read_csv(input_filename)
df['full_path'] = base_dir  + df['path'] + '/' + df['filename']
df.to_csv(input_filename, index=False)

Training

Pretraining ZRNet_mlp

torchrun --nproc_per_node=1 --master_port=4370 basicsr/train_infer.py -opt 'Options/zrnet_mlp.yml' --launcher pytorch

Training ZRNet The pretrained ZRNet_mlp can be downloaded from from here

torchrun --nproc_per_node=1 --master_port=4370 basicsr/train_infer.py -opt 'Options/zrnet_azi_train.yml' --launcher pytorch

Testing

Download the trained models and run the code

torchrun --nproc_per_node=1 --master_port=4370 basicsr/train_infer.py -opt 'Options/zrnet_azi_infer.yml' --launcher pytorch

Notes

You could try training and testing ZRNet with different grouping strategies of Zernike Graphs.

• No grouping: Options/zrnet_nogrp_train.yml
• Aberration grouping: Options/zrnet_ab_train.yml