Improve alignment API, phase mask design, expose waveprop for PSF sim…

…ulation (#130)

* Add option for BGR conversion.

* Add downsample option before Hugging Face upload.

* Remove prints.

* Add option to use waveprop in simulating PSF, improve ROI extraction API.

* Improve load_data interface.

* Add one-shot method for ADMM and GD.

* Refactor is_torch to use_torch.

* Add option to center mask.

* Height map as result for phase mask.

* Make normalize explicit.

* Fix top_left naming.

* Update CHANGELOG.

* Fix mask tests.

CHANGELOG.rst

@@ -36,6 +36,8 @@ Changed
 - For trainable masks, set trainable parameters inside the child class.
 - ``distance_sensor`` optional for ``lensless.hardware.mask.Mask``, e.g. don't need for fabrication.
 - More intuitive interface for MURA for coded aperture (``lensless.hardware.mask.CodedAperture``), i.e. directly pass prime number.
+- ``is_torch`` to ``use_torch`` in ``lensless.hardware.mask.Mask``
+- ``self.height_map`` as characterization of phase masks (instead of phase pattern which can change for each wavelength)
 
 
 Bugfix

configs/benchmark.yaml

@@ -7,15 +7,15 @@ hydra:
     chdir: True
 
 
-dataset: DiffuserCam   # DiffuserCam, DigiCamCelebA, DigiCamHF
+dataset: DiffuserCam   # DiffuserCam, DigiCamCelebA, HFDataset
 seed: 0
 
 huggingface:
   repo: "bezzam/DigiCam-Mirflickr-MultiMask-25K"
   image_res: [900, 1200]  # used during measurement
   rotate: True   # if measurement is upside-down
   alignment:
-    topright: [80, 100]  # height, width
+    top_left: [80, 100]  # height, width
     height: 200
   downsample: 1
 
@@ -88,6 +88,7 @@ simulation:
   scene2mask: 0.25   # [m]
   mask2sensor: 0.002   # [m]     
   # see waveprop.devices
+  use_waveprop: False    # for PSF simulation
   sensor: "rpi_hq"
   snr_db: 10
   # simulate different sensor resolution

configs/train_digicam_multimask.yaml

@@ -25,13 +25,13 @@ files:
       display_res: [900, 1200]  # used during measurement
       rotate: True   # if measurement is upside-down
       alignment:
-        topright: [80, 100]  # height, width
+        top_left: [80, 100]  # height, width
         height: 200
 
 # TODO: these parameters should be in the dataset?
 alignment:
   # when there is no downsampling
-  topright: [80, 100]  # height, width
+  top_left: [80, 100]  # height, width
   height: 200
 
 training:

configs/train_digicam_singlemask.yaml

@@ -24,13 +24,13 @@ files:
       display_res: [900, 1200]  # used during measurement
       rotate: True   # if measurement is upside-down
       alignment:
-        topright: [80, 100]  # height, width
+        top_left: [80, 100]  # height, width
         height: 200
 
 # TODO: these parameters should be in the dataset?
 alignment:
   # when there is no downsampling
-  topright: [80, 100]  # height, width
+  top_left: [80, 100]  # height, width
   height: 200
 
 training:

configs/train_unrolledADMM.yaml

@@ -28,7 +28,7 @@ files:
 
   # -- processing parameters
   downsample: 2    # factor by which to downsample the PSF, note that for DiffuserCam the PSF has 4x the resolution
-  downsample_lensed: 2
+  downsample_lensed: 2   # only used if lensed if measured
   input_snr: null    # adding shot noise at input (for measured dataset) at this SNR in dB
   vertical_shift: null
   horizontal_shift: null
@@ -41,7 +41,7 @@ files:
   extra_eval: null  # dict of extra datasets to evaluate on
 
 alignment: null
-#   topright: null    # height, width
+#   top_left: null    # height, width
   # height: null
 
 torch: True
@@ -129,6 +129,7 @@ simulation:
   scene2mask: 10e-2     # scene2mask: 40e-2       
   mask2sensor: 9e-3     # mask2sensor: 4e-3       
   # see waveprop.devices
+  use_waveprop: False    # for PSF simulation
   sensor: "rpi_hq"
   snr_db: 10
   # simulate different sensor resolution

lensless/eval/benchmark.py

@@ -84,10 +84,6 @@ def benchmark(
         if not os.path.exists(output_dir):
             os.mkdir(output_dir)
 
-    alignment = None
-    if hasattr(dataset, "alignment"):
-        alignment = dataset.alignment
-
     if metrics is None:
         metrics = {
             "MSE": MSELoss().to(device),
@@ -156,13 +152,14 @@ def benchmark(
             prediction = prediction.reshape(-1, *prediction.shape[-3:]).movedim(-1, -3)
             lensed = lensed.reshape(-1, *lensed.shape[-3:]).movedim(-1, -3)
 
-            if alignment is not None:
-                prediction = prediction[
-                    ...,
-                    alignment["topright"][0] : alignment["topright"][0] + alignment["height"],
-                    alignment["topright"][1] : alignment["topright"][1] + alignment["width"],
-                ]
-                # expected that lensed is also reshaped accordingly
+            if hasattr(dataset, "alignment"):
+                if dataset.alignment is not None:
+                    prediction = dataset.extract_roi(prediction, axis=(-2, -1))
+                else:
+                    prediction, lensed = dataset.extract_roi(
+                        prediction, axis=(-2, -1), lensed=lensed
+                    )
+                assert np.all(lensed.shape == prediction.shape)
             elif crop is not None:
                 prediction = prediction[
                     ...,