From 6204ff444651b5e5ddb352665683fbb3574d9db5 Mon Sep 17 00:00:00 2001
From: Eric Bezzam <ebezzam@gmail.com>
Date: Mon, 27 May 2024 12:01:32 +0000
Subject: [PATCH] Use waveprop param.

---
 configs/benchmark.yaml             |  3 ++-
 configs/digicam_example.yaml       | 13 +++++++++++--
 scripts/measure/digicam_example.py | 17 ++++++++++++++++-
 scripts/sim/digicam_psf.py         |  1 -
 4 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/configs/benchmark.yaml b/configs/benchmark.yaml
index 24e0366e..0ae4cce6 100644
--- a/configs/benchmark.yaml
+++ b/configs/benchmark.yaml
@@ -86,7 +86,8 @@ simulation:
   # mask2sensor: 9e-3     # mask2sensor: 4e-3  
   # -- for CelebA
   scene2mask: 0.25   # [m]
-  mask2sensor: 0.002   # [m]     
+  mask2sensor: 0.002   # [m] 
+  deadspace: True    # whether to account for deadspace for programmable mask      
   # see waveprop.devices
   use_waveprop: False    # for PSF simulation
   sensor: "rpi_hq"
diff --git a/configs/digicam_example.yaml b/configs/digicam_example.yaml
index 70af04d7..14eee2fe 100644
--- a/configs/digicam_example.yaml
+++ b/configs/digicam_example.yaml
@@ -10,7 +10,9 @@ rpi:
 # mask parameters
 mask:
   fp: null  # provide path, otherwise generate with seed
-  seed: 1
+  seed: 0
+  # defaults to configuration use for this dataset: https://huggingface.co/datasets/bezzam/DigiCam-Mirflickr-SingleMask-25K
+  # ie this config: configs/collect_mirflickr_singlemask.yaml
   shape: [54, 26]
   center: [57, 77]
 
@@ -37,4 +39,11 @@ capture:
 # reconstruction parameters
 recon:
   torch_device: 'cpu'
-  n_iter: 100    # number of iterations of ADMM
\ No newline at end of file
+  n_iter: 100    # number of iterations of ADMM
+
+simulation:
+  use_waveprop: True
+  deadspace: True
+  scene2mask: 0.3
+  mask2sensor: 0.002
+  gamma: null    # just for saving
\ No newline at end of file
diff --git a/scripts/measure/digicam_example.py b/scripts/measure/digicam_example.py
index f6e51cb6..e75c2436 100644
--- a/scripts/measure/digicam_example.py
+++ b/scripts/measure/digicam_example.py
@@ -19,6 +19,7 @@
 from lensless.utils.io import save_image
 from lensless.hardware.trainable_mask import AdafruitLCD
 from lensless.utils.io import load_image
+from lensless.utils.image import gamma_correction
 
 
 @hydra.main(version_base=None, config_path="../../configs", config_name="digicam_example")
@@ -32,6 +33,7 @@ def digicam(config):
     mask_center = config.mask.center
     torch_device = config.recon.torch_device
     capture_config = config.capture
+    simulation_config = config.simulation
 
     # load mask
     if mask_fp is not None:
@@ -48,14 +50,27 @@ def digicam(config):
         slm="adafruit",
         downsample=capture_config["down"],
         flipud=capture_config["flip"],
+        use_waveprop=simulation_config.get("use_waveprop", False),
+        scene2mask=simulation_config.get("scene2mask", None),
+        mask2sensor=simulation_config.get("mask2sensor", None),
+        deadspace=simulation_config.get("deadspace", True),
         # color_filter=color_filter,
     )
     psf = mask.get_psf().to(torch_device).detach()
     psf_fp = "digicam_psf.png"
-    save_image(psf[0].cpu().numpy(), psf_fp)
+
+    psf_np = psf[0].cpu().numpy()
+    gamma = simulation_config.get("gamma", None)
+    if gamma is not None:
+        psf_np = psf_np / psf_np.max()
+        psf_np = gamma_correction(psf_np, gamma=gamma)
+
+    save_image(psf_np, psf_fp)
     print(f"PSF shape: {psf.shape}")
     print(f"PSF saved to {psf_fp}")
 
+    raise ValueError
+
     if measurement_fp is not None:
         # load image
         img = load_image(
diff --git a/scripts/sim/digicam_psf.py b/scripts/sim/digicam_psf.py
index 1f15e2e0..e4463560 100644
--- a/scripts/sim/digicam_psf.py
+++ b/scripts/sim/digicam_psf.py
@@ -206,7 +206,6 @@ def digicam_psf(config):
         # plot overlayed
         fp = os.path.join(output_folder, "psf_overlay.png")
         psf_meas_norm = psf_meas[0] / np.max(psf_meas)
-        # psf_meas_norm = gamma_correction(psf_meas_norm, gamma=config.digicam.gamma)
         psf_in_np_norm = psf_in_np / np.max(psf_in_np)
 
         plt.figure()