Demo files for icassp and better benchmarking.

ebezzam · ebezzam · commit 71f9680c513f · 2025-04-03T13:24:11.000Z
diff --git a/configs/benchmark/diffusercam.yaml b/configs/benchmark/diffusercam.yaml
@@ -23,22 +23,23 @@ algorithms: [
   # "ADMM",
 
   ## -- reconstructions trained on DiffuserCam measured
-  "hf:diffusercam:mirflickr:U5+Unet8M",
+  # "hf:diffusercam:mirflickr:U5+Unet8M",
   # "hf:diffusercam:mirflickr:Unet8M+U5",
   # "hf:diffusercam:mirflickr:TrainInv+Unet8M",
   # "hf:diffusercam:mirflickr:MMCN4M+Unet4M",
   # "hf:diffusercam:mirflickr:MWDN8M",
-  "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M",
   # "hf:diffusercam:mirflickr:Unet4M+TrainInv+Unet4M",
   # "hf:diffusercam:mirflickr:Unet2M+MMCN+Unet2M",
   # "hf:diffusercam:mirflickr:Unet2M+MWDN6M",
   # "hf:diffusercam:mirflickr:Unet4M+U10+Unet4M",
-  "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
+  # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
 
   # ## comparing UNetRes and Transformer, ADAMW optimizer
   # "hf:diffusercam:mirflickr:Transformer4M+U5+Transformer4M",
   # "hf:diffusercam:mirflickr:Transformer4M+U5+Transformer4M_psfNN",
   # "hf:diffusercam:mirflickr:U5+Transformer8M",
+  "hf:diffusercam:mirflickr:U5+Unet8M_adamw",
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_adamw",
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN_adamw",
 
diff --git a/configs/benchmark/diffusercam_fullres.yaml b/configs/benchmark/diffusercam_fullres.yaml
@@ -23,6 +23,7 @@ algorithms: [
   # "ADMM",
 
   # ## comparing LeADMM5 and SVDeconvNet, ADAMW optimizer
+  "hf:diffusercam:mirflickr:U5+Unet12M_fullres",
   "hf:diffusercam:mirflickr:Unet6M+U5+Unet6M_fullres",
   "hf:diffusercam:mirflickr:Unet6M+U5+Unet6M_psfNN_fullres",
   "hf:diffusercam:mirflickr:SVDecon+UNet8M",
diff --git a/configs/demo_icassp2025.yaml b/configs/demo_icassp2025.yaml
@@ -0,0 +1,37 @@
+defaults:
+  - demo
+  - _self_
+
+  
+gamma: 2.2  # for plotting PSF
+
+display:
+  image_res: [-1,1100]   # negative value to set according to other dimension
+  vshift: -26
+
+capture:
+  legacy: True
+  bayer: True
+  rgb: True
+  down: 4
+  awb_gains: [2.0, 1.1]
+
+recon:
+  gamma: null
+  downsample: 4
+  dtype: float32
+  use_torch: True
+  torch_device: cuda:0
+  algo: admm
+
+  fista:
+    n_iter: 300
+
+  admm:
+    n_iter: 10
+
+postproc:
+  # crop in percent to extract region of interest
+  # set to null to skip
+  crop_hor: [0.3, 0.75]
+  crop_vert: [0.22, 0.7]
diff --git a/configs/recon/diffusercam.yaml b/configs/recon/diffusercam.yaml
@@ -0,0 +1,30 @@
+# python scripts/recon/diffusercam.py
+defaults:
+  - defaults
+  - _self_
+
+cache_dir: null
+idx: 0
+
+# - Learned reconstructions: see "lensless/recon/model_dict.py"
+# --- UNetRes vs Restormer
+# model: U5+Unet8M_adamw
+# model: Unet4M+U5+Unet4M_adamw
+# model: Unet4M+U5+Unet4M_psfNN_adamw
+# model: U5+Transformer8M
+# model: Transformer4M+U5+Transformer4M
+# model: Transformer4M+U5+Transformer4M_psfNN
+# # --- SVDeconv vs LeADMM5
+# model: U5+Unet12M_fullres
+# model: Unet6M+U5+Unet6M_fullres
+# model: Unet6M+U5+Unet6M_psfNN_fullres
+# model: SVDecon+UNet8M
+model: Unet4M+SVDecon+Unet4M
+
+# # -- for ADMM with fixed parameters
+# model: admm
+# n_iter: 100
+
+device: cuda:0
+n_trials: 100   # to get average inference time
+save: True
diff --git a/configs/recon/digicam.yaml b/configs/recon/digicam.yaml
@@ -1,4 +1,4 @@
-# python scripts/recon/digicam_mirflickr.py
+# python scripts/recon/digicam.py
 defaults:
   - defaults
   - _self_
diff --git a/configs/telegram_demo_icassp2025.yaml b/configs/telegram_demo_icassp2025.yaml
@@ -0,0 +1,35 @@
+defaults:
+  - telegram_demo
+  - _self_
+
+# for Telegram
+token: null
+setup_fp: icassp_setup.jpeg
+
+# usernames and IP address
+rpi_username: null
+rpi_hostname: null
+
+config_name: demo_icassp2025
+psf:
+  fp: psf_voronoi_apr2025/voronoi_rgb.png
+  downsample: 4
+
+files:
+  mnist:
+    fp: data/original/mnist_3.png
+    exposure: 0.05
+    brightness: 100
+  thumb: 
+    fp: data/original/thumbs_up.png
+    exposure: 0.02
+    brightness: 80
+  face:
+    fp: data/original/face.jpg
+    exposure: 0.03
+    brightness: 80
+  tree:
+    fp: data/original/tree.png
+    exposure: 0.08
+    brightness: 100
+    
diff --git a/lensless/recon/model_dict.py b/lensless/recon/model_dict.py
@@ -90,13 +90,15 @@
             "Unet4M+U5+Unet4M_ft_tapecam": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-ft-tapecam",
             "Unet4M+U5+Unet4M_ft_tapecam_post": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-ft-tapecam-post",
             "Unet4M+U5+Unet4M_ft_tapecam_pre": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-ft-tapecam-pre",
-            # comparing with transformers, with ADAMW optimizer (rest with ADAM)
+            # comparing with transformers, with ADAMW optimizer (below with ADAMW optimizer)
+            "U5+Unet8M_adamw": "bezzam/diffusercam-mirflickr-unrolled-admm5-unet8M-adamw",
             "Unet4M+U5+Unet4M_adamw": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-adamw",
             "Unet4M+U5+Unet4M_psfNN_adamw": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-psfNN-adamw",
             "U5+Transformer8M": "bezzam/diffusercam-mirflickr-unrolled-admm5-transformer8M",
             "Transformer4M+U5+Transformer4M": "bezzam/diffusercam-mirflickr-transformer4M-unrolled-admm5-transformer4M",
             "Transformer4M+U5+Transformer4M_psfNN": "bezzam/difusercam-mirflickr-transformer4M-unrolled-admm5-transformer4M-psfNN",
             # (~11.6M param) comparing with SVDeconvNet, with ADAMW optimizer (full resolution images)
+            "U5+Unet12M_fullres": "bezzam/diffusercam-mirflickr-unrolled-admm5-unet12M-fullres",
             "Unet6M+U5+Unet6M_fullres": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-fullres",
             "Unet6M+U5+Unet6M_psfNN_fullres": "bezzam/diffusercam-mirflickr-unet4M-unrolled-admm5-unet4M-psfNN-fullres",
             "SVDecon+UNet8M": "bezzam/diffusercam-mirflickr-svdecon-unet4M",
diff --git a/lensless/utils/dataset.py b/lensless/utils/dataset.py
@@ -149,7 +149,7 @@ def print_available_datasets():
         )
 
 
-def get_dataset(dataset_name, split, cache_dir=None, **kwargs):
+def get_dataset(dataset_name, split, **kwargs):
     """
     Get a dataset by name.
 
@@ -159,8 +159,6 @@ def get_dataset(dataset_name, split, cache_dir=None, **kwargs):
         Name of the dataset from the available datasets in ``available_datasets``.
     split : str
         Split of the dataset to load (e.g. "train", "test").
-    cache_dir : str
-        Directory to cache the dataset. By default stored in ~/.cache/huggingface/datasets
 
     Returns
     -------
@@ -177,7 +175,9 @@ def get_dataset(dataset_name, split, cache_dir=None, **kwargs):
     assert split in ["train", "test"], "Split should be 'train' or 'test'"
 
     dataset_config = available_datasets[dataset_name]
-    return HFDataset(split=split, cache_dir=cache_dir, **dataset_config, **kwargs)
+    # replace dataset_config with anything from kwargs
+    dataset_config.update(kwargs)
+    return HFDataset(split=split, **dataset_config)
 
 
 class DualDataset(Dataset):
diff --git a/scripts/measure/analyze_image.py b/scripts/measure/analyze_image.py
@@ -39,6 +39,7 @@
 
 import click
 import cv2
+import os
 import numpy as np
 import matplotlib.pyplot as plt
 from lensless.utils.image import rgb2gray, gamma_correction, resize
@@ -149,14 +150,14 @@ def analyze_image(fp, gamma, width, bayer, lens, lensless, bg, rg, plot_width, s
     ax.set_title("RGB")
     ax = pixel_histogram(img, ax=ax_rgb[1], nbits=nbits)
     ax.set_title("Histogram")
-    fig_rgb.savefig("rgb_analysis.png")
+    fig_rgb.savefig(os.path.join(os.path.dirname(fp), "rgb_analysis.png"))
 
     img_grey = rgb2gray(img[None, ...])
     ax = plot_image(img_grey, gamma=gamma, normalize=True, ax=ax_gray[0])
     ax.set_title("Grayscale")
     ax = pixel_histogram(img_grey, ax=ax_gray[1], nbits=nbits)
     ax.set_title("Histogram")
-    fig_gray.savefig("grey_analysis.png")
+    fig_gray.savefig(os.path.join(os.path.dirname(fp), "grey_analysis.png"))
 
     img_grey = img_grey.squeeze()
     img = img.squeeze()
diff --git a/scripts/measure/prep_display_image.py b/scripts/measure/prep_display_image.py
@@ -98,6 +98,7 @@ def display(
         img = np.zeros((screen_res[1], screen_res[0], 3), dtype=img_og.dtype)
 
         if image_res is None:
+
             # set image with padding and correct aspect ratio
             if screen_res[0] < screen_res[1]:
 
@@ -120,6 +121,20 @@ def display(
             new_height = int(ratio * image_height)
             image_res = (new_width, new_height)
 
+        # if negative value in image res
+        elif image_res[0] < 0 or image_res[1] < 0:
+            assert image_res[0] > 0 or image_res[1] > 0, "Both dimensions cannot be negative."
+            # rescale according to non-negative value
+            if image_res[0] < 0:
+                new_height = image_res[1]
+                ratio = new_height / float(image_width)
+                image_res[0] = int(ratio * image_width)
+
+            elif image_res[1] < 0:
+                new_width = image_res[0]
+                ratio = new_width / float(image_height)
+                image_res[1] = int(ratio * image_height)
+
         # set image within screen
         img_og = cv2.resize(img_og, image_res, interpolation=interpolation)
         img[: image_res[1], : image_res[0]] = img_og
diff --git a/scripts/recon/diffusercam.py b/scripts/recon/diffusercam.py
@@ -0,0 +1,92 @@
+import hydra
+import yaml
+import torch
+from lensless import ADMM
+from lensless.utils.plot import plot_image
+import os
+from lensless.utils.io import save_image
+import time
+from lensless.recon.model_dict import download_model, load_model
+from lensless.utils.dataset import get_dataset
+
+
+@hydra.main(version_base=None, config_path="../../configs/recon", config_name="diffusercam")
+def apply_pretrained(config):
+    save = config.save
+    device = config.device
+    n_trials = config.n_trials
+    model_name = config.model
+    dataset_name = "mirflickr"
+
+    # load config
+    if model_name == "admm":
+        # take config from any model just to get dataset hyperparameters
+        model_path = download_model(camera="diffusercam", dataset=dataset_name, model="U20")
+        config_path = os.path.join(model_path, ".hydra", "config.yaml")
+        with open(config_path, "r") as stream:
+            model_config = yaml.safe_load(stream)
+
+    else:
+        model_path = download_model(camera="diffusercam", dataset=dataset_name, model=model_name)
+        config_path = os.path.join(model_path, ".hydra", "config.yaml")
+        with open(config_path, "r") as stream:
+            model_config = yaml.safe_load(stream)
+    model_config["files"]["cache_dir"] = config.cache_dir
+
+    # load data
+    test_set = get_dataset(
+        dataset_name="diffusercam_mirflickr",
+        split="test",
+        **model_config["files"],
+    )
+    psf = test_set.psf.to(device)
+    print("Test set size: ", len(test_set))
+
+    # reconstruction specified example
+    idx = config.idx
+    lensless, lensed = test_set[idx]
+    lensless = lensless.to(device)
+
+    # load model
+    if model_name == "admm":
+        recon = ADMM(psf, n_iter=config.n_iter)
+    else:
+        # load best model
+        recon = load_model(model_path, psf, device)
+
+    # print data shape
+    print(f"Data shape :  {lensless.shape}")
+
+    # apply reconstruction
+    start_time = time.time()
+    for _ in range(n_trials):
+        with torch.no_grad():
+
+            recon.set_data(lensless)
+            res = recon.apply(
+                disp_iter=-1,
+                save=False,
+                gamma=None,
+                plot=False,
+            )
+    end_time = time.time()
+    avg_time_ms = (end_time - start_time) / n_trials * 1000
+    print(f"Avg inference [ms] : {avg_time_ms} ms")
+
+    img = res[0].cpu().numpy().squeeze()
+
+    plot_image(img)
+    if lensed is not None:
+        plot_image(lensed)
+
+    if save:
+        print(f"Saving images to {os.getcwd()}")
+        save_image(img, f"{model_name}_{idx}.png")
+        if lensed is not None:
+            lensed_np = lensed[0].cpu().numpy()
+            save_image(lensed_np, f"original_{idx}.png")
+        save_image(lensless[0].cpu().numpy(), f"lensless_{idx}.png")
+
+
+if __name__ == "__main__":
+    apply_pretrained()

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-# python scripts/recon/digicam_mirflickr.py`
	`1`	`+# python scripts/recon/digicam.py`
`2`	`2`	`defaults:`
`3`	`3`	`- defaults`
`4`	`4`	`- _self_`