From a79279ef6d0ec15dc8c074ae728bd07b1fc454c3 Mon Sep 17 00:00:00 2001 From: Eric Bezzam Date: Sat, 25 May 2024 20:15:56 +0200 Subject: [PATCH] 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 | 3 +- configs/defaults_recon.yaml | 1 + configs/demo.yaml | 1 + configs/{digicam.yaml => digicam_config.yaml} | 0 configs/digicam_example.yaml | 40 ++++++ lensless/hardware/utils.py | 6 +- lensless/utils/io.py | 25 +++- scripts/hardware/config_digicam.py | 2 +- scripts/hardware/digicam_measure_psfs.py | 2 +- scripts/hardware/set_digicam_mask_distance.py | 2 +- scripts/measure/digicam_example.py | 122 ++++++++++++++++++ scripts/measure/remote_capture.py | 33 +---- scripts/recon/admm.py | 3 + 13 files changed, 203 insertions(+), 37 deletions(-) rename configs/{digicam.yaml => digicam_config.yaml} (100%) create mode 100644 configs/digicam_example.yaml create mode 100644 scripts/measure/digicam_example.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 7825c43f..7e4f9caf 100644 --- a/CHANGELOG.rst +++ b/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 ~~~~~~~ diff --git a/configs/defaults_recon.yaml b/configs/defaults_recon.yaml index 08d783ac..97613b50 100644 --- a/configs/defaults_recon.yaml +++ b/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' diff --git a/configs/demo.yaml b/configs/demo.yaml index c4dfee47..2f6c29ac 100644 --- a/configs/demo.yaml +++ b/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 diff --git a/configs/digicam.yaml b/configs/digicam_config.yaml similarity index 100% rename from configs/digicam.yaml rename to configs/digicam_config.yaml diff --git a/configs/digicam_example.yaml b/configs/digicam_example.yaml new file mode 100644 index 00000000..70af04d7 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/lensless/hardware/utils.py b/lensless/hardware/utils.py index f409d7cb..7777f201 100644 --- a/lensless/hardware/utils.py +++ b/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( diff --git a/lensless/utils/io.py b/lensless/utils/io.py index 6d07bc27..7e44975b 100644 --- a/lensless/utils/io.py +++ b/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) diff --git a/scripts/hardware/config_digicam.py b/scripts/hardware/config_digicam.py index 0807519b..4c64bdf3 100644 --- a/scripts/hardware/config_digicam.py +++ b/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 diff --git a/scripts/hardware/digicam_measure_psfs.py b/scripts/hardware/digicam_measure_psfs.py index 901d24cb..5085c91d 100644 --- a/scripts/hardware/digicam_measure_psfs.py +++ b/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 diff --git a/scripts/hardware/set_digicam_mask_distance.py b/scripts/hardware/set_digicam_mask_distance.py index dcd0dd79..e9360353 100644 --- a/scripts/hardware/set_digicam_mask_distance.py +++ b/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 diff --git a/scripts/measure/digicam_example.py b/scripts/measure/digicam_example.py new file mode 100644 index 00000000..f6e51cb6 --- /dev/null +++ b/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() diff --git a/scripts/measure/remote_capture.py b/scripts/measure/remote_capture.py index b7b5c6c0..a38f722a 100644 --- a/scripts/measure/remote_capture.py +++ b/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) diff --git a/scripts/recon/admm.py b/scripts/recon/admm.py index 12584e59..1d1c261c 100644 --- a/scripts/recon/admm.py +++ b/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"],