Add full digicam example. (#119)

* Add full digicam example.

* Update digicam config.

* Clean up digicam example.

* Setting relative path.

* Set image resolution remotely.

* Add checks that are in on-device script.

* Add options for background image.

* Update CHANGELOG.

CHANGELOG.rst

@@ -13,7 +13,8 @@ Unreleased
 Added
 ~~~~~
 
-- Nothing
+- Option to pass background image to ``utils.io.load_data``.
+- Option to set image resolution with ``hardware.utils.display`` function.
 
 Changed
 ~~~~~~~

configs/defaults_recon.yaml

@@ -9,6 +9,7 @@ input:
   data: data/raw_data/thumbs_up_rgb.png
   dtype: float32
   original: null   # ground truth image
+  background: null   # background image
 
 torch: False
 torch_device: 'cpu'

configs/demo.yaml

@@ -16,6 +16,7 @@ rpi:
 display:
   # default to this screen: https://www.dell.com/en-us/work/shop/dell-ultrasharp-usb-c-hub-monitor-u2421e/apd/210-axmg/monitors-monitor-accessories#techspecs_section
   screen_res: [1920, 1200]   # width, height
+  image_res: null
   pad: 0
   hshift: 0
   vshift: -10

configs/digicam_example.yaml

@@ -0,0 +1,40 @@
+# python scripts/measure/digicam_example.py
+hydra:
+  job:
+    chdir: True    # change to output folder
+
+rpi:
+  username: null
+  hostname: null
+
+# mask parameters
+mask:
+  fp: null  # provide path, otherwise generate with seed
+  seed: 1
+  shape: [54, 26]
+  center: [57, 77]
+
+# measurement parameters
+capture:
+  fp: null
+  exp: 0.5
+  sensor: rpi_hq
+  script: ~/LenslessPiCam/scripts/measure/on_device_capture.py
+  iso: 100
+  config_pause: 1
+  sensor_mode: "0"
+  nbits_out: 8
+  nbits_capture: 12
+  legacy: True
+  gray: False
+  fn: raw_data
+  bayer: True
+  awb_gains: [1.6, 1.2]
+  rgb: True
+  down: 8
+  flip: True
+
+# reconstruction parameters
+recon:
+  torch_device: 'cpu'
+  n_iter: 100    # number of iterations of ADMM

lensless/hardware/utils.py

@@ -199,8 +199,7 @@ def capture(
         if verbose:
             print(f"\nCopying over picture as {localfile}...")
         os.system(
-            'scp "%s@%s:%s" %s >%s'
-            % (rpi_username, rpi_hostname, remotefile, localfile, NULL_FILE)
+            'scp "%s@%s:%s" %s >%s' % (rpi_username, rpi_hostname, remotefile, localfile, NULL_FILE)
         )
 
         if rgb or gray:
@@ -242,6 +241,7 @@ def display(
     rpi_username,
     rpi_hostname,
     screen_res,
+    image_res=None,
     brightness=100,
     rot90=0,
     pad=0,
@@ -279,6 +279,8 @@ def display(
     prep_command = f"{rpi_python} {script} --fp {remote_tmp_file} \
         --pad {pad} --vshift {vshift} --hshift {hshift} --screen_res {screen_res[0]} {screen_res[1]} \
         --brightness {brightness} --rot90 {rot90} --output_path {display_path} "
+    if image_res is not None:
+        prep_command += f" --image_res {image_res[0]} {image_res[1]}"
     if verbose:
         print(f"COMMAND : {prep_command}")
     subprocess.Popen(

lensless/utils/io.py

@@ -378,6 +378,7 @@ def load_psf(
 def load_data(
     psf_fp,
     data_fp,
+    background_fp=None,
     return_float=True,
     downsample=None,
     bg_pix=(5, 25),
@@ -495,10 +496,32 @@ def load_data(
         as_4d=True,
         return_float=return_float,
         shape=shape,
-        normalize=normalize,
+        normalize=normalize if background_fp is None else False,
         bgr_input=bgr_input,
     )
 
+    if background_fp is not None:
+        bg = load_image(
+            background_fp,
+            flip=flip,
+            bayer=bayer,
+            blue_gain=blue_gain,
+            red_gain=red_gain,
+            as_4d=True,
+            return_float=return_float,
+            shape=shape,
+            normalize=False,
+            bgr_input=bgr_input,
+        )
+        assert bg.shape == data.shape
+
+        data -= bg
+        # clip to 0
+        data = np.clip(data, a_min=0, a_max=data.max())
+
+        if normalize:
+            data /= data.max()
+
     if data.shape != psf.shape:
         # in DiffuserCam dataset, images are already reshaped
         data = resize(data, shape=psf.shape)

scripts/hardware/config_digicam.py

@@ -11,7 +11,7 @@
 from lensless.hardware.utils import set_mask_sensor_distance
 
 
-@hydra.main(version_base=None, config_path="../../configs", config_name="digicam")
+@hydra.main(version_base=None, config_path="../../configs", config_name="digicam_config")
 def config_digicam(config):
 
     rpi_username = config.rpi.username

scripts/hardware/digicam_measure_psfs.py

@@ -8,7 +8,7 @@
 SATURATION_THRESHOLD = 0.01
 
 
-@hydra.main(version_base=None, config_path="../../configs", config_name="digicam")
+@hydra.main(version_base=None, config_path="../../configs", config_name="digicam_config")
 def config_digicam(config):
 
     rpi_username = config.rpi.username

scripts/hardware/set_digicam_mask_distance.py

@@ -2,7 +2,7 @@
 from lensless.hardware.utils import set_mask_sensor_distance
 
 
-@hydra.main(version_base=None, config_path="../../configs", config_name="digicam")
+@hydra.main(version_base=None, config_path="../../configs", config_name="digicam_config")
 def config_digicam(config):
 
     rpi_username = config.rpi.username

scripts/measure/digicam_example.py

@@ -0,0 +1,122 @@
+"""
+DigiCam example to remotely:
+1. Set mask pattern.
+2. Capture image.
+3. Reconstruct image with simulated PSF.
+
+TODO: display image. At the moment should be done with `scripts/measure/remote_display.py`
+
+"""
+
+
+import hydra
+from hydra.utils import to_absolute_path
+import numpy as np
+from lensless.hardware.slm import set_programmable_mask, adafruit_sub2full
+from lensless.hardware.utils import capture
+import torch
+from lensless import ADMM
+from lensless.utils.io import save_image
+from lensless.hardware.trainable_mask import AdafruitLCD
+from lensless.utils.io import load_image
+
+
+@hydra.main(version_base=None, config_path="../../configs", config_name="digicam_example")
+def digicam(config):
+    measurement_fp = config.capture.fp
+    mask_fp = config.mask.fp
+    seed = config.mask.seed
+    rpi_username = config.rpi.username
+    rpi_hostname = config.rpi.hostname
+    mask_shape = config.mask.shape
+    mask_center = config.mask.center
+    torch_device = config.recon.torch_device
+    capture_config = config.capture
+
+    # load mask
+    if mask_fp is not None:
+        mask_vals = np.load(to_absolute_path(mask_fp))
+    else:
+        # create random mask within [0, 1]
+        np.random.seed(seed)
+        mask_vals = np.random.uniform(0, 1, mask_shape)
+
+    # simulate PSF
+    mask = AdafruitLCD(
+        initial_vals=torch.from_numpy(mask_vals.astype(np.float32)),
+        sensor=capture_config["sensor"],
+        slm="adafruit",
+        downsample=capture_config["down"],
+        flipud=capture_config["flip"],
+        # 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)
+    print(f"PSF shape: {psf.shape}")
+    print(f"PSF saved to {psf_fp}")
+
+    if measurement_fp is not None:
+        # load image
+        img = load_image(
+            to_absolute_path(measurement_fp),
+            verbose=True,
+        )
+
+    else:
+        ## measure data
+        # -- prepare full mask
+        pattern = adafruit_sub2full(
+            mask_vals,
+            center=mask_center,
+        )
+
+        # -- set mask
+        print("Setting mask")
+        set_programmable_mask(
+            pattern,
+            "adafruit",
+            rpi_username=rpi_username,
+            rpi_hostname=rpi_hostname,
+        )
+
+        # -- capture
+        print("Capturing")
+        localfile, img = capture(
+            rpi_username=rpi_username,
+            rpi_hostname=rpi_hostname,
+            verbose=False,
+            **capture_config,
+        )
+        print(f"Captured to {localfile}")
+
+    """ analyze image """
+    print("image range: ", img.min(), img.max())
+
+    """ reconstruction """
+    # -- normalize
+    img = img.astype(np.float32) / img.max()
+    # prep
+    img = torch.from_numpy(img)
+    # -- if [H, W, C] -> [D, H, W, C]
+    if len(img.shape) == 3:
+        img = img.unsqueeze(0)
+    if capture_config["flip"]:
+        img = torch.rot90(img, dims=(-3, -2), k=2)
+
+    # reconstruct
+    print("Reconstructing")
+    recon = ADMM(psf)
+    recon.set_data(img.to(psf.device))
+    res = recon.apply(disp_iter=None, plot=False, n_iter=config.recon.n_iter)
+    res_np = res[0].cpu().numpy()
+    res_np = res_np / res_np.max()
+    lensless_np = img[0].cpu().numpy()
+    save_image(lensless_np, "digicam_raw.png")
+    save_image(res_np, "digicam_recon.png")
+
+    print("Done")
+
+
+if __name__ == "__main__":
+    digicam()

scripts/measure/remote_capture.py

@@ -32,7 +32,6 @@
 import cv2
 from pprint import pprint
 import matplotlib.pyplot as plt
-import rawpy
 from lensless.hardware.utils import check_username_hostname
 from lensless.hardware.sensor import SensorOptions, sensor_dict, SensorParam
 from lensless.utils.image import rgb2gray, print_image_info
@@ -127,9 +126,12 @@ def liveview(config):
         and "bullseye" in result_dict["RPi distribution"]
         and not legacy
     ):
+        assert not rgb or not gray, "RGB and gray not supported for RPi HQ sensor"
 
         if bayer:
 
+            assert config.capture.down is None
+
             # copy over DNG file
             remotefile = f"~/{remote_fn}.dng"
             localfile = os.path.join(save, f"{fn}.dng")
@@ -138,35 +140,6 @@ def liveview(config):
 
             img = load_image(localfile, verbose=True, bayer=bayer, nbits_out=nbits_out)
 
-            # raw = rawpy.imread(localfile)
-
-            # # https://letmaik.github.io/rawpy/api/rawpy.Params.html#rawpy.Params
-            # # https://www.libraw.org/docs/API-datastruct-eng.html
-            # if nbits_out > 8:
-            #     # only 8 or 16 bit supported by postprocess
-            #     if nbits_out != 16:
-            #         print("casting to 16 bit...")
-            #     output_bps = 16
-            # else:
-            #     if nbits_out != 8:
-            #         print("casting to 8 bit...")
-            #     output_bps = 8
-            # img = raw.postprocess(
-            #     adjust_maximum_thr=0,  # default 0.75
-            #     no_auto_scale=False,
-            #     # no_auto_scale=True,
-            #     gamma=(1, 1),
-            #     output_bps=output_bps,
-            #     bright=1,  # default 1
-            #     exp_shift=1,
-            #     no_auto_bright=True,
-            #     # use_camera_wb=True,
-            #     # use_auto_wb=False,
-            #     # -- gives better balance for PSF measurement
-            #     use_camera_wb=False,
-            #     use_auto_wb=True,  # default is False? f both use_camera_wb and use_auto_wb are True, then use_auto_wb has priority.
-            # )
-
             # print image properties
             print_image_info(img)
 

scripts/recon/admm.py

@@ -29,6 +29,9 @@ def admm(config):
     psf, data = load_data(
         psf_fp=to_absolute_path(config.input.psf),
         data_fp=to_absolute_path(config.input.data),
+        background_fp=to_absolute_path(config.input.background)
+        if config.input.background is not None
+        else None,
         dtype=config.input.dtype,
         downsample=config["preprocess"]["downsample"],
         bayer=config["preprocess"]["bayer"],