Measuring datasets with 3D printed masks (#124)

* Update setup.

* Add option to reconstruct during measurement.

* Add FZA measurement config.

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.
 

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:

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

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,

rpi_requirements.txt

@@ -5,4 +5,5 @@ matplotlib>=3.4.2
 hydra-code
 paramiko
 numpy>=1.24.2
-scipy>=1.6.0
+scipy>=1.6.0
+git+https://github.com/ebezzam/slm-controller.git

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