Improve plotting for demo notebook and add config for ICASSP demo.

ebezzam · web-flow · commit 6c3e1b23eb7f · 2025-04-13T19:40:13.000+02:00
Update demos files
diff --git a/.github/workflows/python_no_pycsou.yml b/.github/workflows/python_no_pycsou.yml
@@ -42,7 +42,7 @@ jobs:
       run: |
         pip install flake8
         # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=B,C,E,F,W,T4,B9 --show-source --statistics --max-complexity=18 --max-line-length=100 --ignore=E203,E266,E501,W503,F403,F401,C901 --exclude=profile,scripts,docs
+        flake8 . --count --select=B,C,E,F,W,T4,B9 --show-source --statistics --max-complexity=18 --max-line-length=100 --ignore=E203,E266,E501,W503,F403,F401,C901,PLW0602 --exclude=profile,scripts,docs
         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
         # flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Format with black
diff --git a/.github/workflows/python_pycsou.yml b/.github/workflows/python_pycsou.yml
@@ -42,7 +42,7 @@ jobs:
       run: |
         pip install flake8
         # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=B,C,E,F,W,T4,B9 --show-source --statistics --max-complexity=18 --max-line-length=100 --ignore=E203,E266,E501,W503,F403,F401,C901 --exclude=profile,scripts,docs
+        flake8 . --count --select=B,C,E,F,W,T4,B9 --show-source --statistics --max-complexity=18 --max-line-length=100 --ignore=E203,E266,E501,W503,F403,F401,C901,PLW0602 --exclude=profile,scripts,docs
         # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
         # flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Format with black
diff --git a/configs/benchmark/diffusercam.yaml b/configs/benchmark/diffusercam.yaml
@@ -20,25 +20,29 @@ huggingface:
   single_channel_psf: True
 
 algorithms: [
-  # "ADMM",
+  "ADMM",
+  "ADMM_PnP",
+  "FISTA",
+  "FISTA_PnP",
 
   ## -- 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",
 
@@ -83,6 +87,8 @@ algorithms: [
   # "hf:digicam:mirflickr_multi_25k:Unet4M+U5+Unet4M_wave",
 ]
 
-save_idx: [0, 1, 3, 4, 8, 45, 58, 63]
-n_iter_range: [100]    # for ADMM
+save_idx: [0, 1, 3, 4, 8]
+n_iter_range: [10]    # for ADMM
+# save_idx: [0, 1, 3, 4, 8, 45, 58, 63]
+# n_iter_range: [100]    # for ADMM
 
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/benchmark/digicam.yaml b/configs/benchmark/digicam.yaml
@@ -23,19 +23,21 @@ huggingface:
 
 algorithms: [
   # "ADMM",
+  "ADMM",
+  "ADMM_PnP",
 
   # # -- reconstructions trained on measured data
-  "hf:digicam:mirflickr_single_25k:U5+Unet8M_wave",
+  # "hf:digicam:mirflickr_single_25k:U5+Unet8M_wave",
   # "hf:digicam:mirflickr_single_25k:Unet8M+U5_wave",
   # "hf:digicam:mirflickr_single_25k:TrainInv+Unet8M_wave",
   # "hf:digicam:mirflickr_single_25k:MMCN4M+Unet4M_wave",
   # "hf:digicam:mirflickr_single_25k:MWDN8M_wave",
   # "hf:digicam:mirflickr_single_25k:Unet4M+TrainInv+Unet4M_wave",
-  "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave",
+  # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave",
   # "hf:digicam:mirflickr_single_25k:Unet2M+MMCN+Unet2M_wave",
   # "hf:digicam:mirflickr_single_25k:Unet2M+MWDN6M_wave",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U10+Unet4M_wave",
-  "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
+  # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_psfNN",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_flips",
   # "hf:digicam:mirflickr_single_25k:Unet4M+U5+Unet4M_wave_flips_rotate10",
   # "hf:digicam:mirflickr_single_25k:Unet8M_wave_v2",
@@ -53,8 +55,10 @@ algorithms: [
   # "hf:diffusercam:mirflickr:Unet4M+U5+Unet4M_psfNN",
 ]
 
-save_idx: [1, 2, 4, 5, 9, 24, 33, 61]
-n_iter_range: [100]    # for ADMM
+# save_idx: [1, 2, 4, 5, 9, 24, 33, 61]
+# n_iter_range: [100]    # for ADMM
+save_idx: [1, 2, 4, 5, 9]
+n_iter_range: [10]    # for ADMM
 
 # simulating PSF
 simulation:
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: [900,-1]   # negative value to set according to other dimension
+  vshift: -26
+
+capture:
+  legacy: True
+  bayer: True
+  rgb: True
+  down: 4
+  awb_gains: [1.8, 1.1]
+
+recon:
+  gamma: null
+  downsample: 4
+  dtype: float32
+  use_torch: True
+  torch_device: cuda:0
+  algo: fista
+
+  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.29, 0.67]
+  crop_vert: [0.22, 0.78]
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,51 @@
+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
+
+default_algo: fista
+
+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
+
+overlay:
+  alpha: 60
+
+  img1:
+    fp: data/original/epfl_white.png
+    scaling: 0.23
+    position: [5,15]
+
+  img2:
+    fp: data/original/icassp.png
+    scaling: 0.2
+    position: [305,10]
+
+  img3: null
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/recon/utils.py b/lensless/recon/utils.py
@@ -307,11 +307,12 @@ def apply_denoiser(
             noise_level = noise_level / 255.0
         else:
             noise_level = torch.tensor([noise_level / 255.0])
-
         image = torch.cat(
             (
                 image,
-                noise_level.repeat(image.shape[0], 1, image.shape[2], image.shape[3]),
+                noise_level.repeat(image.shape[0], 1, image.shape[2], image.shape[3]).to(
+                    image.device
+                ),
             ),
             dim=1,
         )
@@ -633,7 +634,7 @@ def __init__(
             Epoch at which to unfreeze post process component. Default is None.
 
         """
-        global print
+        # global print
 
         self.use_wandb = use_wandb
 
diff --git a/lensless/utils/dataset.py b/lensless/utils/dataset.py
@@ -101,7 +101,7 @@ def natural_sort(arr):
     # DigiCam MirFlickr Mini (472 MB) https://huggingface.co/datasets/bezzam/DigiCam-Mirflickr-SingleMask-1K
     "digicam_mirflickr_mini": {
         "size (GB)": 0.472,
-        "huggingface_repo": "bezzam/DigiCam-Mirflickr-SingleMask-25K",
+        "huggingface_repo": "bezzam/DigiCam-Mirflickr-SingleMask-1K",
         "display_res": [900, 1200],
         "rotate": True,
         "alignment": {"top_left": [80, 100], "height": 200},
@@ -138,6 +138,29 @@ def natural_sort(arr):
         "display_res": [600, 600],
         "alignment": {"top_left": [118, 220], "height": 123},
     },
+    # Multilens MirFlickr Mini (427 MB) https://huggingface.co/datasets/Lensless/mirflickr_voronoi_1k
+    "multilens_mirflickr_mini": {
+        "size (GB)": 0.427,
+        "huggingface_repo": "Lensless/mirflickr_voronoi_1k",
+        "psf": "psf_measured.png",
+        "display_res": [900, 1200],
+    },
+    # Coded Aperture (MLS) MirFlickr 1K (467 MB) https://huggingface.co/datasets/Lensless/mirflickr_CA_fine_1k
+    "mls_mirflickr_1k": {
+        "size (GB)": 0.467,
+        "huggingface_repo": "Lensless/mirflickr_CA_fine_1k",
+        "psf": "psf_measured.png",
+        "display_res": [900, 1200],
+        # "alignment": {"top_left": [118, 220], "height": 123},
+    },
+    # Fresnel Zone Aperture MirFlickr 1K (454 MB) https://huggingface.co/datasets/Lensless/Mirflickr_FZA_fine_1k
+    "fza_mirflickr_1k": {
+        "size (GB)": 0.454,
+        "huggingface_repo": "Lensless/Mirflickr_FZA_fine_1k",
+        "psf": "psf_measured.png",
+        "display_res": [900, 1200],
+        # "alignment": {"top_left": [118, 220], "height": 123},
+    },
 }
 
 
@@ -149,7 +172,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 +182,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 +198,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/lensless/utils/plot.py b/lensless/utils/plot.py
diff --git a/notebook/lenslesspicam_demo.ipynb b/notebook/lenslesspicam_demo.ipynb
diff --git a/scripts/demo/telegram_bot.py b/scripts/demo/telegram_bot.py
diff --git a/scripts/measure/analyze_image.py b/scripts/measure/analyze_image.py
diff --git a/scripts/measure/prep_display_image.py b/scripts/measure/prep_display_image.py
diff --git a/scripts/recon/diffusercam.py b/scripts/recon/diffusercam.py

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_`