Merge branch 'main' into global-shutter

.github/workflows/python_pycsou.yml

@@ -59,5 +59,5 @@ jobs:
         pip install -U pytest
         pip install -r recon_requirements.txt
         pip install -r mask_requirements.txt
-        pip install git+https://github.com/matthieumeo/pycsou.git@v2-dev
+        pip install git+https://github.com/matthieumeo/pycsou.git@38e9929c29509d350a7ff12c514e2880fdc99d6e
         pytest

CHANGELOG.rst

@@ -13,6 +13,29 @@ Unreleased
 Added
 ~~~~~
 
+- Trainable reconstruction can return intermediate outputs (between pre- and post-processing).
+- Auto-download for DRUNet model.
+- ``utils.dataset.DiffuserCamMirflickr`` helper class for Mirflickr dataset.
+
+Changed
+~~~~~~~
+
+- Better logic for saving best model. Based on desired metric rather than last epoch, and intermediate models can be saved.
+- Optional normalization in ``utils.io.load_image``.
+
+Bugfix
+~~~~~~
+
+- Support for unrolled reconstruction with grayscale, needed to copy to three channels for LPIPS.
+- Fix bad train/test split for DiffuserCamMirflickr in unrolled training.
+
+
+1.0.5 - (2023-09-05)
+--------------------
+
+Added
+~~~~~
+
 - Sensor module.
 - Single-script and Telegram demo.
 - Link and citation for JOSS.
@@ -22,8 +45,15 @@ Added
 - Script for measuring arbitrary dataset (from Raspberry Pi).
 - Support for preprocessing and postprocessing, such as denoising, in ``TrainableReconstructionAlgorithm``. Both trainable and fix postprocessing can be used.
 - Utilities to load a trained DruNet model for use as postprocessing in ``TrainableReconstructionAlgorithm``.
+- Unified interface for dataset. See ``utils.dataset.DualDataset``.
+- New simulated dataset compatible with new data format ([(batch_size), depth, width, height, color]). See ``utils.dataset.SimulatedFarFieldDataset``.
+- New dataset for pair of original image and their measurement from a screen. See ``utils.dataset.MeasuredDataset`` and ``utils.dataset.MeasuredDatasetSimulatedOriginal``.
 - Support for unrolled loading and inference in the script ``admm.py``.
-- Tikhonov reconstruction for coded aperture measurements (MLS / MURA).
+- Tikhonov reconstruction for coded aperture measurements (MLS / MURA): numpy and Pytorch support.
+- New ``Trainer`` class to train ``TrainableReconstructionAlgorithm`` with PyTorch.
+- New ``TrainableMask`` and ``TrainablePSF`` class to train/fine-tune a mask from a dataset.
+- New ``SimulatedDatasetTrainableMask`` class to train/fine-tune a mask for measurement.
+- PyTorch support for ``lensless.utils.io.rgb2gray``.
 
 
 Changed

README.rst

@@ -60,7 +60,8 @@ Python 3.9, as some Python library versions may not be available with
 earlier versions of Python. Moreover, its `end-of-life <https://endoflife.date/python>`__ 
 is Oct 2025.
 
-**Local machine**
+*Local machine setup*
+=====================
 
 Below are commands that worked for our configuration (Ubuntu
 21.04), but there are certainly other ways to download a repository and
@@ -83,16 +84,20 @@ install the library locally.
    # (optional) try reconstruction on local machine
    python scripts/recon/admm.py
 
+   # (optional) try reconstruction on local machine with GPU
+   python scripts/recon/admm.py -cn pytorch
 
-Note (25-04-2023): for using reconstruction method based on Pycsou ``lensless.apgd.APGD``,
-V2 has to be installed:
+
+Note (25-04-2023): for using the :py:class:`~lensless.recon.apgd.APGD` reconstruction method based on Pycsou
+(now `Pyxu <https://github.com/matthieumeo/pyxu>`__), a specific commit has 
+to be installed (as there was no release at the time of implementation):
 
 .. code:: bash
 
-   pip install git+https://github.com/matthieumeo/pycsou.git@v2-dev
+   pip install git+https://github.com/matthieumeo/pycsou.git@38e9929c29509d350a7ff12c514e2880fdc99d6e
 
 If PyTorch is installed, you will need to be sure to have PyTorch 2.0 or higher, 
-as Pycsou V2 is not compatible with earlier versions of PyTorch. Moreover, 
+as Pycsou is not compatible with earlier versions of PyTorch. Moreover, 
 Pycsou requires Python within 
 `[3.9, 3.11) <https://github.com/matthieumeo/pycsou/blob/v2-dev/setup.cfg#L28>`__.
 
@@ -102,7 +107,8 @@ Moreover, ``numba`` (requirement for Pycsou V2) may require an older version of
 
    pip install numpy==1.23.5
 
-**Raspberry Pi**
+*Raspberry Pi setup*
+====================
 
 After `flashing your Raspberry Pi with SSH enabled <https://medium.com/@bezzam/setting-up-a-raspberry-pi-without-a-monitor-headless-9a3c2337f329>`__, 
 you need to set it up for `passwordless access <https://medium.com/@bezzam/headless-and-passwordless-interfacing-with-a-raspberry-pi-ssh-453dd75154c3>`__. 

configs/defaults_recon.yaml

@@ -8,11 +8,13 @@ input:
   # File path for raw data  
   data: data/raw_data/thumbs_up_rgb.png
   dtype: float32
+  original: null   # ground truth image
 
 torch: False
 torch_device: 'cpu'
 
 preprocess:
+  normalize: True
   # Downsampling factor along X and Y
   downsample: 4
   # Image shape (height, width) for reconstruction.

configs/diffusercam_mirflickr_single_admm.yaml

@@ -0,0 +1,43 @@
+# python scripts/recon/admm.py -cn diffusercam_mirflickr_single_admm
+defaults:
+  - defaults_recon
+  - _self_
+
+
+display:
+  gamma: null
+
+input:
+  # File path for recorded PSF
+  psf: data/DiffuserCam_Test/psf.tiff
+  # File path for raw data  
+  data: data/DiffuserCam_Test/diffuser/im5.npy
+  dtype: float32
+  original: data/DiffuserCam_Test/lensed/im5.npy
+
+torch: True
+torch_device: 'cuda:0'
+
+preprocess:
+  downsample: 8   # factor for PSF, which is 4x resolution of image
+  normalize: False
+
+admm:
+  # Number of iterations
+  n_iter: 20
+  # Hyperparameters
+  mu1: 1e-6
+  mu2: 1e-5
+  mu3: 4e-5
+  tau: 0.0001
+  #Loading unrolled model
+  unrolled: True
+  # checkpoint_fp: pretrained_models/Pre_Unrolled_Post-DiffuserCam/model_weights.pt
+  checkpoint_fp: outputs/2023-09-11/22-06-49/recon.pt  # pre unet and post drunet
+  pre_process_model: 
+    network : UnetRes  # UnetRes or DruNet or null
+    depth : 2 # depth of each up/downsampling layer. Ignore if network is DruNet
+  post_process_model: 
+    network : DruNet  # UnetRes or DruNet or null
+    depth : 2 # depth of each up/downsampling layer. Ignore if network is DruNet
+

configs/fine-tune_PSF.yaml

@@ -0,0 +1,18 @@
+# python scripts/recon/train_unrolled.py -cn fine-tune_PSF
+defaults:
+  - train_unrolledADMM
+  - _self_
+
+#Trainable Mask
+trainable_mask:
+  mask_type: TrainablePSF #Null or "TrainablePSF"
+  initial_value: psf
+  mask_lr: 1e-3
+  L1_strength: 1.0  #False or float
+
+#Training
+training:
+  save_every: 5
+
+display:
+  gamma: 2.2

configs/mask_sim_single.yaml

@@ -8,6 +8,7 @@ files:
   #original: data/original/mnist_3.png
 
 save: True
+use_torch: False
 
 simulation:
   object_height: 0.3

configs/train_pre-post-processing.yaml

@@ -0,0 +1,24 @@
+# python scripts/recon/train_unrolled.py -cn train_pre-post-processing
+defaults:
+  - train_unrolledADMM
+  - _self_
+
+display:
+  disp: 400
+
+reconstruction:
+  method: unrolled_admm
+
+  pre_process:
+    network: UnetRes
+    depth: 2
+  post_process:
+    network: DruNet
+    depth: 4
+
+training:
+  epoch: 50
+  slow_start: 0.01
+
+loss: l2
+lpips: 1.0

configs/train_psf_from_scratch.yaml

@@ -0,0 +1,18 @@
+# python scripts/recon/train_unrolled.py -cn train_psf_from_scratch
+defaults:
+  - train_unrolledADMM
+  - _self_
+
+# Train Dataset
+files:
+  dataset: mnist   # Simulated : "mnist", "fashion_mnist", "cifar10", "CelebA". Measure :"DiffuserCam"
+  celeba_root: /scratch/bezzam
+  downsample: 8
+
+#Trainable Mask
+trainable_mask:
+  mask_type: TrainablePSF #Null or "TrainablePSF"
+  initial_value: "random"
+
+simulation:
+  grayscale: False

configs/unrolled_recon.yaml → configs/train_unrolledADMM.yaml

@@ -1,33 +1,24 @@
+# python scripts/recon/train_unrolled.py
 hydra:
   job:
     chdir: True    # change to output folder
 
-#Reconstruction algorithm
-input:
-  # File path for recorded PSF
-  psf: data/DiffuserCam_Mirflickr_200_3011302021_11h43_seed11/psf.tiff
-  dtype: float32
+# Dataset
+files:
+  dataset: data/DiffuserCam   # Simulated : "mnist", "fashion_mnist", "cifar10", "CelebA". Measure :"DiffuserCam"
+  celeba_root: null   # path to parent directory of CelebA: https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
+  psf: data/psf.tiff
+  diffusercam_psf: True
+  n_files: null    # null to use all for both train/test
+  downsample: 2    # factor by which to downsample the PSF, note that for DiffuserCam the PSF has 4x the resolution
 
 torch: True
 torch_device: 'cuda'
 
-preprocess:
-  # Image shape (height, width) for reconstruction.
-  shape: null
-  # Whether image is raw bayer data.
-  bayer: False
-  blue_gain: null
-  red_gain: null
-  # Same PSF for all channels (sum) or unique PSF for RGB.
-  single_psf: False
-  # Whether to perform construction in grayscale.
-  gray: False
-
-
 display:
   # How many iterations to wait for intermediate plot. 
   # Set to negative value for no intermediate plots.
-  disp: 400
+  disp: 500
   # Whether to plot results.
   plot: True
   # Gamma factor for plotting.
@@ -48,63 +39,70 @@ reconstruction:
     learn_tk: True
   unrolled_admm:
     # Number of iterations
-    n_iter: 5
+    n_iter: 20
     # Hyperparameters
     mu1: 1e-4
     mu2: 1e-4
     mu3: 1e-4
     tau: 2e-4
   pre_process: 
-    network : UnetRes  # UnetRes or DruNet or null
+    network : null  # UnetRes or DruNet or null
     depth : 2 # depth of each up/downsampling layer. Ignore if network is DruNet
   post_process: 
-    network : UnetRes  # UnetRes or DruNet or null
+    network : null  # UnetRes or DruNet or null
     depth : 2 # depth of each up/downsampling layer. Ignore if network is DruNet
 
-# Train Dataset
-
-files:
-  dataset: "DiffuserCam"   # "mnist", "fashion_mnist", "cifar10", "CelebA", "DiffuserCam"
-  n_files: null    # null to use all
+#Trainable Mask
+trainable_mask:
+  mask_type: Null #Null or "TrainablePSF"
+  # "random" (with shape of config.files.psf) or "psf" (using config.files.psf)
+  initial_value: psf
+  grayscale: False
+  mask_lr: 1e-3
+  L1_strength: 1.0  #False or float
 
 target: "object_plane"    # "original" or "object_plane" or "label"
 
 #for simulated dataset
 simulation:
   grayscale: False
   # random variations
-  object_height: 0.6   # range for random height or scalar
+  object_height: 0.04   # range for random height or scalar
+  flip: True # change the orientation of the object (from vertical to horizontal)
   random_shift: False
   random_vflip: 0.5
   random_hflip: 0.5
   random_rotate: False
   # these distance parameters are typically fixed for a given PSF
-  # for tape_rgb psf      # for DiffuserCam psf
-  scene2mask: 40e-2       # scene2mask: 10e-2
-  mask2sensor: 4e-3       # mask2sensor: 9e-3
+  # for DiffuserCam psf # for tape_rgb psf     
+  scene2mask: 10e-2     # scene2mask: 40e-2       
+  mask2sensor: 9e-3     # mask2sensor: 4e-3       
   # see waveprop.devices
   sensor: "rpi_hq"
-  snr_db: 40
+  snr_db: 10
   # simulate different sensor resolution
   # output_dim: [24, 32]    # [H, W] or null
   # Downsampling for PSF
   downsample: 8
   # max val in simulated measured (quantized 8 bits)
+  quantize: False   # must be False for differentiability
   max_val: 255
 
 #Training
 
 training:
   batch_size: 8
   epoch: 50
+  metric_for_best_model: null   # e.g. LPIPS_Vgg, null does test loss
+  save_every: null
   #In case of instable training
   skip_NAN: True
   slow_start: False  #float how much to reduce lr for first epoch
 
 
 optimizer:
   type: Adam
-  lr: 1e-6
+  lr: 1e-4
 
 loss: 'l2'
 # set lpips to false to deactivate. Otherwise, give the weigth for the loss (the main loss l2/l1 always having a weigth of 1)

docs/requirements.txt

@@ -4,4 +4,6 @@ docutils==0.16     # >0.17 doesn't render bullets
 numpy>=1.22      # so that default dtype are correctly rendered
 torch>=1.10
 torchvision>=0.15.2
-torchmetrics>=0.11.4
+torchmetrics>=0.11.4
+pyFFS>=2.2.3   # for waveprop
+waveprop>=0.0.7

docs/source/conf.py

@@ -21,21 +21,25 @@
     "torchmetrics.image",
     "scipy.ndimage",
     "pycsou.abc",
+    "pycsou.operator",
     "pycsou.operator.func",
+    "pycsou.operator.linop",
     "pycsou.opt.solver",
     "pycsou.opt.stop",
     "pycsou.runtime",
     "pycsou.util",
     "pycsou.util.ptype",
     "PIL",
+    "PIL.Image",
     "tqdm",
     "paramiko",
     "paramiko.ssh_exception",
     "perlin_numpy",
-    "waveprop",
-    "waveprop.fresnel",
-    "waveprop.rs",
-    "waveprop.noise",
+    "hydra",
+    "hydra.utils",
+    "scipy.special",
+    "matplotlib.cm",
+    "pyffs",
 ]
 for mod_name in MOCK_MODULES:
     sys.modules[mod_name] = mock.Mock()