Add instructions for training.

Author: ebezzam

configs/train/README.md

@@ -0,0 +1,87 @@
+# Training configurations
+
+The commands below train with the corresponding dataset, which will be downloaded from Hugging Face.
+
+By commenting/uncommenting relevant sections in the configuration file, you can train the models with different architectures. By default, the model architecture uses five unrolleed iterations of ADMM for camera inversion, and UNetRes models for the pre-processor post-processor, and PSF correction.
+
+With DiffuserCam, we show how to set different camera inversion methods.
+
+## DiffuserCam
+
+The commands below show how to train different camera inversion methods on the DiffuserCam dataset (downsampled by a factor of 2 along each dimension).For a fair comparison, all models use around 8.1M parameters.
+
+```bash
+# unrolled ADMM
+python scripts/recon/train_learning_based.py -cn diffusercam
+
+# Trainable inversion (FlatNet but with out adversarial loss)
+# -- need to set PSF as trainable
+python scripts/recon/train_learning_based.py -cn diffusercam \
+    reconstruction.method=trainable_inv \
+    reconstruction.psf_network=False \
+    trainable_mask.mask_type=TrainablePSF \
+	trainable_mask.L1_strength=False
+
+# Unrolled ADMM with compensation branch
+# - adjust shapes of pre and post processors
+python scripts/recon/train_learning_based.py -cn diffusercam \
+    reconstruction.psf_network=False \
+    reconstruction.pre_process.nc=[16,32,64,128] \
+    reconstruction.post_process.nc=[16,32,64,128] \
+    reconstruction.compensation=[24,64,128,256,400]
+
+# Multi wiener deconvolution network (MWDN) 
+# with PSF correction built into the network
+python scripts/recon/train_learning_based.py -cn diffusercam \
+    reconstruction.method=multi_wiener \
+    reconstruction.multi_wiener.nc=[32,64,128,256,436] \
+    reconstruction.pre_process.network=null \
+    reconstruction.post_process.network=null \
+    reconstruction.psf_network=False
+```
+
+Similar to [PhoCoLens](https://phocolens.github.io/), we can train a camera inversion model that learns multiple PSFs. The training below uses the DiffuserCam dataset with its full resolution, and the number of model parameters is around 11.6M.
+```bash
+python scripts/recon/train_learning_based.py -cn diffusercam \
+    reconstruction.method=svdeconvnet \
+    reconstruction.pre_process.nc=[32,64,116,128] \
+    reconstruction.psf_network=False \
+    trainable_mask.mask_type=TrainablePSF \
+	trainable_mask.L1_strength=False \
+    files.downsample=1 files.downsample_lensed=1
+```
+
+## TapeCam
+
+```bash
+# unrolled ADMM
+python scripts/recon/train_learning_based.py -cn tapecam
+```
+
+## DigiCam (Single Mask)
+
+```bash
+# unrolled ADMM
+python scripts/recon/train_learning_based.py -cn digicam
+```
+
+## DigiCam (Multiple Masks)
+
+```bash
+# unrolled ADMM
+python scripts/recon/train_learning_based.py -cn digicam_multimask
+```
+
+## DigiCam CelebA
+
+```bash
+# unrolled ADMM
+python scripts/recon/train_learning_based.py -cn digicam_celeba
+```
+
+## MultiPSF under External Illumination
+
+```bash
+# unrolled ADMM
+python scripts/recon/train_learning_based.py -cn multilens_ambient
+```

lensless/recon/multi_wiener.py

@@ -279,7 +279,7 @@ def apply(self, **kwargs):
         # apply to data
         return self.forward(self._data, **kwargs)
 
-    def reconstruction_error(self, prediction, lensless):
+    def reconstruction_error(self, prediction, lensless, **kwargs):
         convolver = self._convolver
         if not convolver.pad:
             prediction = convolver._pad(prediction)

scripts/recon/train_learning_based.py

@@ -670,8 +670,8 @@ def train_learned(config):
                 psf_channels = 3
 
             assert config.reconstruction.direct_background_subtraction is False, "Not supported"
-            assert config.reconstruction.learned_background_subtraction is None, "Not supported"
-            assert config.reconstruction.integrated_background_subtraction is None, "Not supported"
+            assert config.reconstruction.learned_background_subtraction is False, "Not supported"
+            assert config.reconstruction.integrated_background_subtraction is False, "Not supported"
             assert psf_network is None, "Not supported"
 
             recon = MultiWiener(