PSF compensation network and adding noise to input/PSF for robustness…

… tests. (#158)

* Fixes for when no background is present.

* Update setup.

* Add support for authentication for learned models.

* Fix PSF used for reconstruction debug examples.

* Add script for psf err vs metric.

* Add option for setting SNR with digital noise.

* Add support for parameterize and perturb.

* Update experiments.

* Update macos in GitHub actions.

* Update gitignore.

* Update.

* Update CHANGELOG.

.github/workflows/python_pycsou.yml

@@ -20,7 +20,7 @@ jobs:
       fail-fast: false
       max-parallel: 12
       matrix:
-        os: [ubuntu-latest, macos-12, windows-latest]
+        os: [ubuntu-latest, macos-latest, windows-latest]
         python-version: [3.9, "3.10"]
     steps:
     - uses: actions/checkout@v3

.gitignore

@@ -12,6 +12,7 @@ notebooks/models/*
 *.png
 *.jpg
 *.npy
+*.sh
 
 configs/telegram_demo_secret.yaml
 

CHANGELOG.rst

@@ -34,6 +34,9 @@ Added
 - Integrated background subtraction object ``lensless.recon.integrated_background.IntegratedBackgroundSub``.
 - Option to concatenate background to input to pre-processor.
 - Add support for datasets with measured background to ``lensless.utils.dataset.HFDataset``.
+- Add parameterize and peturb to evaluate model adaptation.
+- PSF correction network.
+- Option to add noise to input image or PSF (for robustness experiments).
 
 
 Changed
@@ -48,6 +51,7 @@ Bugfix
 - Computation of average metric in batches.
 - Support for grayscale PSF for RealFFTConvolve2D.
 - Calling model.eval() before inference, and model.train() before training.
+- Normalization when computing reconstruction error.
 
 
 1.0.7 - (2024-05-14)
@@ -65,7 +69,6 @@ Added
 - Support for training/testing with multiple mask patterns in the dataset.
 - Multi-GPU support for training.
 - Dataset which interfaces with Hugging Face (``lensless.utils.dataset.HFDataset``).
-- Scripts for authentication.
 - DigiCam support for Telegram demo.
 - DiffuserCamMirflickr Hugging Face API.
 - Fallback for normalization if data not in 8bit range (``lensless.utils.io.save_image``).

README.rst

@@ -193,7 +193,8 @@ A few students at EPFL have also contributed to this project:
 * Julien Sahli: support and extension of algorithms for 3D.
 * Yohann Perron: unrolled algorithms for reconstruction.
 * Aaron Fargeon: mask designs.
-* Rein Bentdal: mask fabrication with 3D printing.
+* Rein Bentdal and David Karoubi: mask fabrication with 3D printing.
+* Stefan Peters: imaging under external illumination.
 
 Citing this work
 ----------------

configs/benchmark.yaml

@@ -10,6 +10,7 @@ hydra:
 dataset: DiffuserCam   # DiffuserCam, DigiCamCelebA, HFDataset
 seed: 0
 batchsize: 1    # must be 1 for iterative approaches
+snr: null
 
 huggingface:
   repo: "bezzam/DigiCam-Mirflickr-MultiMask-25K"
@@ -42,6 +43,8 @@ algorithms: ["ADMM", "ADMM_Monakhova2019", "FISTA"] #["ADMM", "ADMM_Monakhova201
 baseline: "MONAKHOVA 100iter"
 
 save_idx: [0, 1, 2, 3, 4]   # provide index of files to save e.g. [1, 5, 10]
+save_intermediate: False   # save intermediate results, i.e. after pre-processor and after camera inversion
+swap_channels: False   # list of two RGB channels to swap, e.g. [0, 1] for swapping red and green
 gamma_psf: 1.5    # gamma factor for PSF
 
 
@@ -57,6 +60,14 @@ admm:
   mu3: 4e-5
   tau: 0.0001
 
+# parameterize and perturb (Gilton et. al.)
+pnp: null
+# pnp:
+#   mu: 1e-3   # weight for distance from original model parameters
+#   lr: 3e-3   # learning rate for SGD
+#   n_iter: 100   # number of iterations
+#   model_path: null    # leave null to be overwritten in script
+
 
 # for DigiCamCelebA
 files:

configs/benchmark_diffusercam_mirflickr.yaml

@@ -20,10 +20,11 @@ huggingface:
   single_channel_psf: True
 
 algorithms: [
-  "ADMM",
+  # "ADMM",
 
   ## -- reconstructions trained on DiffuserCam measured
   "hf:diffusercam:mirflickr:U5+Unet8M",
+  "hf:diffusercam:mirflickr:Unet8M+U5",
   "hf:diffusercam:mirflickr:TrainInv+Unet8M",
   "hf:diffusercam:mirflickr:MMCN4M+Unet4M",
   "hf:diffusercam:mirflickr:MWDN8M",
@@ -32,6 +33,21 @@ algorithms: [
   "hf:diffusercam:mirflickr:Unet2M+MMCN+Unet2M",
   "hf:diffusercam:mirflickr:Unet2M+MWDN6M",
   "hf:diffusercam:mirflickr:Unet4M+U10+Unet4M",
+  "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
+
+  # # -- benchmark PSF error
+  # "hf:diffusercam:mirflickr:U5+Unet8M_psf0dB",
+  # "hf:diffusercam:mirflickr:U5+Unet8M_psf-5dB",
+  # "hf:diffusercam:mirflickr:U5+Unet8M_psf-10dB",
+  # "hf:diffusercam:mirflickr:U5+Unet8M_psf-20dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psf-0dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psf-5dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psf-10dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psf-20dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN_psf-0dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN_psf-10dB",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN_psf-20dB",
+
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_ft_tapecam",
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_ft_tapecam_post",
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_ft_tapecam_pre",
@@ -50,7 +66,9 @@ algorithms: [
   # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave",
   # "hf:digicam:celeba_26k:Unet4M+U5+Unet4M_wave",
   # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M",
+  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_psfNN",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U10+Unet4M_wave",
+  # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
   # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips",
   # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips_rotate10",
   # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_aux1",

configs/benchmark_digicam_celeba.yaml

@@ -6,20 +6,22 @@ defaults:
 
 dataset: HFDataset
 batchsize: 10
-device: "cuda:0"
+device: "cuda:1"
 
 algorithms: [
-  "ADMM",
+  # "ADMM",
 
   ## -- reconstructions trained on measured data
   "hf:digicam:celeba_26k:U5+Unet8M_wave",
+  "hf:digicam:celeba_26k:Unet8M+U5_wave",
   "hf:digicam:celeba_26k:TrainInv+Unet8M_wave",
   "hf:digicam:celeba_26k:MWDN8M_wave",
   "hf:digicam:celeba_26k:MMCN4M+Unet4M_wave",
   "hf:digicam:celeba_26k:Unet2M+MWDN6M_wave",
   "hf:digicam:celeba_26k:Unet4M+TrainInv+Unet4M_wave",
   "hf:digicam:celeba_26k:Unet2M+MMCN+Unet2M_wave",
   "hf:digicam:celeba_26k:Unet4M+U5+Unet4M_wave",
+  "hf:digicam:celeba_26k:Unet4M+U5+Unet4M_wave_psfNN",
   "hf:digicam:celeba_26k:Unet4M+U10+Unet4M_wave",
 
   # # -- reconstructions trained on other datasets/systems
@@ -37,6 +39,7 @@ n_iter_range: [100]    # for ADMM
 huggingface:
   repo: bezzam/DigiCam-CelebA-26K
   psf: psf_simulated_waveprop.png    # psf_simulated_waveprop.png, psf_simulated.png, psf_measured.png
+  cache_dir: /dev/shm
   split_seed: 0
   test_size: 0.15
   downsample: 2

configs/benchmark_digicam_mirflickr_multi.yaml

@@ -25,11 +25,14 @@ algorithms: [
 
   ## -- reconstructions trained on measured data 
   "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave",
+  "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave_psfNN",
   "hf:digicam:mirflickr_multi_25k:Unet4M+U10+Unet4M_wave",
   "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave_aux1",
   "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave_flips",
+  "hf:digicam:mirflickr_multi_25k:Unet8M_wave_v2",
 
   # ## -- reconstructions trained on other datasets/systems
+  # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
   # "hf:diffusercam:mirflickr:Unet4M+U10+Unet4M",
   # "hf:tapecam:mirflickr:Unet4M+U10+Unet4M",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U10+Unet4M_wave",
@@ -42,6 +45,9 @@ algorithms: [
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_ft_flips",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_ft_flips_rotate10",
+  # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
+  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_psfNN",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
 ]
 
 # # -- to only use output from unrolled

configs/benchmark_digicam_mirflickr_pnp.yaml

@@ -0,0 +1,44 @@
+# python scripts/eval/benchmark_recon.py -cn benchmark_digicam_mirflickr_pnp
+defaults:
+  - benchmark
+  - _self_
+
+
+dataset: HFDataset
+batchsize: 1
+device: "cuda:0"
+
+huggingface:
+  repo: "bezzam/DigiCam-Mirflickr-MultiMask-25K"
+  psf: null   # null for simulating PSF
+  image_res: [900, 1200]  # used during measurement
+  rotate: True   # if measurement is upside-down
+  flipud: False
+  flip_lensed: False   # if rotate or flipud is True, apply to lensed
+  alignment:
+    top_left: [80, 100]  # height, width
+    height: 200
+  downsample: 1
+
+algorithms: [
+  "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
+  "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave_psfNN",
+]
+
+pnp:
+  mu: 1e-3   # weight for distance from original model parameters
+  model_path: null    # leave null to be overwritten in script
+  # lr: 1e-2
+  # n_iter: 10
+  lr: 3e-3   # learning rate for SGD
+  n_iter: 10   # number of iterations
+
+
+save_idx: [1, 2, 4, 5, 9, 24, 33, 61]
+
+# simulating PSF
+simulation:
+  use_waveprop: True
+  deadspace: True
+  scene2mask: 0.3
+  mask2sensor: 0.002

configs/benchmark_digicam_mirflickr_single.yaml

@@ -9,7 +9,7 @@ device: "cuda:0"
 
 huggingface:
   repo: "bezzam/DigiCam-Mirflickr-SingleMask-25K"
-  cache_dir: /dev/shm
+  cache_dir: null
   psf: null   # null for simulating PSF
   image_res: [900, 1200]  # used during measurement
   rotate: True   # if measurement is upside-down
@@ -22,10 +22,11 @@ huggingface:
 
 
 algorithms: [
-  "ADMM",
+  # "ADMM",
 
-  # -- reconstructions trained on measured data
+  # # -- reconstructions trained on measured data
   "hf:digicam:mirflickr_single_25k:U5+Unet8M_wave",
+  "hf:digicam:mirflickr_single_25k:Unet8M+U5_wave",
   "hf:digicam:mirflickr_single_25k:TrainInv+Unet8M_wave",
   "hf:digicam:mirflickr_single_25k:MMCN4M+Unet4M_wave",
   "hf:digicam:mirflickr_single_25k:MWDN8M_wave",
@@ -34,8 +35,10 @@ algorithms: [
   "hf:digicam:mirflickr_single_25k:Unet2M+MMCN+Unet2M_wave",
   "hf:digicam:mirflickr_single_25k:Unet2M+MWDN6M_wave",
   "hf:digicam:mirflickr_single_25k:Unet4M+U10+Unet4M_wave",
+  "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
   "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_flips",
   "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_flips_rotate10",
+  "hf:digicam:mirflickr_single_25k:Unet8M_wave_v2",
 
   # ## -- reconstructions trained on other datasets/systems
   # "hf:diffusercam:mirflickr:Unet4M+U10+Unet4M",
@@ -45,9 +48,12 @@ algorithms: [
   # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M",
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M",
   # "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave",
+  # "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave_psfNN",
+  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_psfNN",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
 ]
 
-save_idx: [1, 2, 4, 5, 9]
+save_idx: [1, 2, 4, 5, 9, 24, 33, 61]
 n_iter_range: [100]    # for ADMM
 
 # simulating PSF

configs/benchmark_tapecam_mirflickr.yaml

@@ -5,11 +5,11 @@ defaults:
 
 dataset: HFDataset
 batchsize: 4
-device: "cuda:0"
+device: "cuda:1"
 
 huggingface:
   repo: "bezzam/TapeCam-Mirflickr-25K"
-  cache_dir: /dev/shm
+  cache_dir: null #/dev/shm
   psf: psf.png
   image_res: [900, 1200]  # used during measurement
   rotate: False   # if measurement is upside-down
@@ -23,20 +23,22 @@ huggingface:
 
 ## -- reconstructions trained with same dataset/system
 algorithms: [
-  "ADMM",
+  # "ADMM",
 
   # -- reconstructions trained on measured data
   "hf:tapecam:mirflickr:U5+Unet8M",
+  "hf:tapecam:mirflickr:Unet8M+U5",
   "hf:tapecam:mirflickr:TrainInv+Unet8M",
   "hf:tapecam:mirflickr:MMCN4M+Unet4M",
   "hf:tapecam:mirflickr:Unet4M+U5+Unet4M",
   "hf:tapecam:mirflickr:Unet4M+TrainInv+Unet4M",
   "hf:tapecam:mirflickr:Unet2M+MMCN+Unet2M",
   "hf:tapecam:mirflickr:Unet4M+U10+Unet4M",
-  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips_rotate10",
-  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_aux1",
-  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips",
-  # "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips_rotate10",
+  "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips_rotate10",
+  "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_aux1",
+  "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips",
+  "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_flips_rotate10",
+  "hf:tapecam:mirflickr:Unet4M+U5+Unet4M_psfNN",
 
   # # below models need `single_channel_psf = True`
   # "hf:tapecam:mirflickr:MWDN8M",    
@@ -58,7 +60,9 @@ algorithms: [
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_ft_tapecam_pre",
   # "hf:diffusercam:mirflickr_sim:Unet4M+U5+Unet4M_ft_digicam_multi_pre",
   # "hf:diffusercam:mirflickr_sim:Unet4M+U5+Unet4M_ft_digicam_multi",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
+  # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
 ]
 
 save_idx: [1, 2, 4, 5, 9]
-n_iter_range: [100]    # for ADM
+n_iter_range: [100]    # for ADMM