From 17462ff4cdf476b35f28440acb79c92dc9fdf9c6 Mon Sep 17 00:00:00 2001
From: Eric Bezzam <ebezzam@gmail.com>
Date: Mon, 27 May 2024 16:43:22 +0200
Subject: [PATCH] Measuring datasets with 3D printed masks (#124)

* Update setup.

* Add option to reconstruct during measurement.

* Add FZA measurement config.
---
 README.rst                                   |  3 +-
 configs/collect_dataset.yaml                 |  1 +
 configs/collect_mirflickr_fza.yaml           | 32 +++++++++++++++++++
 lensless/recon/recon.py                      |  4 +--
 rpi_requirements.txt                         |  3 +-
 scripts/measure/collect_dataset_on_device.py | 33 ++++++++++++++++++++
 6 files changed, 71 insertions(+), 5 deletions(-)
 create mode 100644 configs/collect_mirflickr_fza.yaml

diff --git a/README.rst b/README.rst
index 532e0940..428fd07e 100644
--- a/README.rst
+++ b/README.rst
@@ -168,8 +168,7 @@ directory):
    pip install -r rpi_requirements.txt
 
    # test on-device camera capture (after setting up the camera)
-   source lensless_env/bin/activate
-   python scripts/measure/on_device_capture.py
+   (lensless_env) python scripts/measure/on_device_capture.py
 
 You may still need to manually install ``numpy`` and/or ``scipy`` with ``pip`` in case libraries (e.g. ``libopenblas.so.0``) cannot be detected.
    
diff --git a/configs/collect_dataset.yaml b/configs/collect_dataset.yaml
index c898ab2b..c5f56d45 100644
--- a/configs/collect_dataset.yaml
+++ b/configs/collect_dataset.yaml
@@ -24,6 +24,7 @@ min_level: 200
 max_tries: 6
 
 masks: null   # for multi-mask measurements
+recon: null   # parameters for reconstruction (for debugging purposes, not recommended to do during actual measurement as it will significantly increase the time)
 
 # -- display parameters
 display:
diff --git a/configs/collect_mirflickr_fza.yaml b/configs/collect_mirflickr_fza.yaml
new file mode 100644
index 00000000..bd6de24f
--- /dev/null
+++ b/configs/collect_mirflickr_fza.yaml
@@ -0,0 +1,32 @@
+# python scripts/measure/collect_dataset_on_device.py -cn collect_mirflickr_fza
+defaults:
+  - collect_dataset
+  - _self_
+
+input_dir: /mnt/mirflickr/all
+min_level: 170
+
+# FOR TESTING PURPOSE
+output_dir: data/fza_test   # RPi won't have enough memory for full measured dataset
+max_tries: 0
+recon:
+  psf: data/psf/tape_rgb.png  # TODO: set correct PSF
+  n_iter: 10
+# # FOR FINAL MEASUREMENT
+# max_tries: 2
+# output_dir: /mnt/mirflickr/fza_10K
+
+# files to measure
+n_files: 25000
+
+# -- display parameters
+display:
+  image_res: [900, 1200]
+  vshift: -26
+  brightness: 90
+  delay: 1
+
+capture:
+  down: 8
+  exposure: 0.7
+  awb_gains: [1.6, 1.2]  # red, blue, TODO for your mask
\ No newline at end of file
diff --git a/lensless/recon/recon.py b/lensless/recon/recon.py
index 39f78a0c..77813b0e 100644
--- a/lensless/recon/recon.py
+++ b/lensless/recon/recon.py
@@ -493,9 +493,9 @@ def _get_numpy_data(self, data):
     def apply(
         self,
         n_iter=None,
-        disp_iter=10,
+        disp_iter=-1,
         plot_pause=0.2,
-        plot=True,
+        plot=False,
         save=False,
         gamma=None,
         ax=None,
diff --git a/rpi_requirements.txt b/rpi_requirements.txt
index 6c7bf14b..01ccda68 100644
--- a/rpi_requirements.txt
+++ b/rpi_requirements.txt
@@ -5,4 +5,5 @@ matplotlib>=3.4.2
 hydra-code
 paramiko
 numpy>=1.24.2
-scipy>=1.6.0
\ No newline at end of file
+scipy>=1.6.0
+git+https://github.com/ebezzam/slm-controller.git
\ No newline at end of file
diff --git a/scripts/measure/collect_dataset_on_device.py b/scripts/measure/collect_dataset_on_device.py
index 76393b9c..69bd65d9 100644
--- a/scripts/measure/collect_dataset_on_device.py
+++ b/scripts/measure/collect_dataset_on_device.py
@@ -89,6 +89,22 @@ def collect_dataset(config):
                 mask_vals = np.random.uniform(0, 1, config.masks.shape)
                 np.save(mask_fp, mask_vals)
 
+    recon = None
+    if config.recon is not None:
+        # initialize ADMM reconstruction
+        from lensless import ADMM
+        from lensless.utils.io import load_psf
+
+        psf, bg = load_psf(
+            fp=config.recon.psf,
+            downsample=config.capture.down,   # assume full resolution PSF
+            return_bg=True
+        )
+        recon = ADMM(psf, n_iter=config.recon.n_iter)
+
+        recon_dir = plib.Path(output_dir) / "recon"
+        recon_dir.mkdir(exist_ok=True)
+
     # assert input directory exists
     assert os.path.exists(input_dir)
 
@@ -315,6 +331,23 @@ def collect_dataset(config):
                     brightness_vals.append(current_screen_brightness)
                     n_tries_vals.append(n_tries)
 
+        if recon is not None:
+
+            # normalize and remove background
+            output = output.astype(np.float32)
+            output /= output.max()
+            output -= bg
+            output = np.clip(output, a_min=0, a_max=output.max())
+
+            # set data
+            output = output[np.newaxis, :, :, :]
+            recon.set_data(output)
+
+            # reconstruct and save
+            res = recon.apply()
+            recon_fp = recon_dir / output_fp.name
+            save_image(res, recon_fp)
+
         # check if runtime is exceeded
         if config.runtime:
             proc_time = time.time() - start_time