Add demo notebook.

LCAV · Nov 7, 2024 · 80881e1 · 80881e1
1 parent 99d6341
commit 80881e1
Show file tree

Hide file tree

Showing 4 changed files with 1,064 additions and 2 deletions.
diff --git a/lensless/hardware/mask.py b/lensless/hardware/mask.py
@@ -772,7 +772,7 @@ def create_mask(self):
         self.height_map = height_map
 
 
-def phase_retrieval(target_psf, wv, d1, dz, n=1.2, n_iter=10, height_map=False):
+def phase_retrieval(target_psf, wv, d1, dz, n=1.2, n_iter=10, height_map=False, phase_wrap=1):
     """
     Iterative phase retrieval algorithm similar to `PhlatCam <https://ieeexplore.ieee.org/document/9076617>`_,
     using Fresnel propagation.
@@ -791,7 +791,10 @@ def phase_retrieval(target_psf, wv, d1, dz, n=1.2, n_iter=10, height_map=False):
         Refractive index of the mask substrate. Default is 1.2.
     n_iter: int
         Number of iterations. Default value is 10.
+    phase_wrap : int
+        How many multiple of (2*pi) to wrap the phase, to limit heights. Default is 1.
     """
+    assert isinstance(phase_wrap, int), "phase_wrap should be an integer"
     M_p = np.sqrt(target_psf)
 
     if hasattr(d1, "__len__"):
@@ -809,7 +812,7 @@ def phase_retrieval(target_psf, wv, d1, dz, n=1.2, n_iter=10, height_map=False):
         # constrain amplitude to be sqrt(PSF)
         M_p = np.sqrt(target_psf) * np.exp(1j * np.angle(M_p))
 
-    phi = (np.angle(M_phi) + 2 * np.pi) % (2 * np.pi)
+    phi = (np.angle(M_phi) + 2 * np.pi) % (2 * np.pi * phase_wrap)
 
     if height_map:
         return phi, wv * phi / (2 * np.pi * (n - 1))

diff --git a/lensless/utils/dataset.py b/lensless/utils/dataset.py
@@ -47,6 +47,140 @@ def natural_sort(arr):
     return sorted(arr, key=alphanum_key)
 
 
+# available datasets
+available_datasets = {
+    # -- DiffuserCam MirFlickr (7.58 GB) https://huggingface.co/datasets/bezzam/DiffuserCam-Lensless-Mirflickr-Dataset-NORM
+    "diffusercam_mirflickr": {
+        "size (GB)": 7.58,
+        "huggingface_repo": "bezzam/DiffuserCam-Lensless-Mirflickr-Dataset-NORM",
+        "psf": "psf.tiff",
+        "single_channel_psf": True,
+        "flipud": True,
+        "flip_lensed": True,
+        "downsample": 2,
+        "downsample_lensed": 2,
+    },
+    # -- TapeCam MirFlickr (10.5 GB) https://huggingface.co/datasets/bezzam/TapeCam-Mirflickr-25K
+    "tapecam_mirflickr": {
+        "size (GB)": 10.5,
+        "huggingface_repo": "bezzam/TapeCam-Mirflickr-25K",
+        "psf": "psf.png",
+        "display_res": [900, 1200],
+        "alignment": {"top_left": [45, 95], "height": 250},
+    },
+    # -- DigiCam CelebA (33.9 GB) https://huggingface.co/datasets/bezzam/DigiCam-CelebA-26K
+    "digicam_celeba": {
+        "size (GB)": 33.9,
+        "huggingface_repo": "bezzam/DigiCam-CelebA-26K",
+        "psf": "psf_simulated.png",
+        "rotate": True,
+        "split_seed": 0,
+        "downsample": 2,
+        "alignment": {"crop": {"vertical": [0, 525], "horizontal": [265, 695]}},
+        "simulation": {
+            "scene2mask": 0.25,
+            "mask2sensor": 0.002,
+            "object_height": 0.33,
+            "sensor": "rpi_hq",
+            "snr_db": None,
+            "downsample": None,
+            "random_vflip": False,
+            "random_hflip": False,
+            "quantize": False,
+            "vertical_shift": -117,
+            "horizontal_shift": -25,
+        },
+    },
+    # -- DigiCam MirFlickr (11.9 GB) https://huggingface.co/datasets/bezzam/DigiCam-Mirflickr-SingleMask-25K
+    "digicam_mirflickr": {
+        "size (GB)": 11.9,
+        "huggingface_repo": "bezzam/DigiCam-Mirflickr-SingleMask-25K",
+        "display_res": [900, 1200],
+        "rotate": True,
+        "alignment": {"top_left": [80, 100], "height": 200},
+    },
+    # 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",
+        "display_res": [900, 1200],
+        "rotate": True,
+        "alignment": {"top_left": [80, 100], "height": 200},
+    },
+    # -- DigiCam MirFlickr Multimask (12 GB) https://huggingface.co/datasets/bezzam/DigiCam-Mirflickr-MultiMask-25K
+    "digicam_mirflickr_multi": {
+        "size (GB)": 12,
+        "huggingface_repo": "bezzam/DigiCam-Mirflickr-MultiMask-25K",
+        "display_res": [900, 1200],
+        "rotate": True,
+        "alignment": {"top_left": [80, 100], "height": 200},
+    },
+    # -- DigiCam MirFlickr Multimask Mini (477 MB) https://huggingface.co/datasets/bezzam/DigiCam-Mirflickr-MultiMask-1K
+    "digicam_mirflickr_multi_mini": {
+        "size (GB)": 0.477,
+        "huggingface_repo": "bezzam/DigiCam-Mirflickr-MultiMask-25K",
+        "display_res": [900, 1200],
+        "rotate": True,
+        "alignment": {"top_left": [80, 100], "height": 200},
+    },
+    # MultiLens MirFlickr Ambient (16.7 GB) https://huggingface.co/datasets/Lensless/MultiLens-Mirflickr-Ambient
+    "multilens_mirflickr_ambient": {
+        "size (GB)": 16.7,
+        "huggingface_repo": "Lensless/MultiLens-Mirflickr-Ambient",
+        "psf": "psf.png",
+        "display_res": [600, 600],
+        "alignment": {"top_left": [118, 220], "height": 123},
+    },
+    # MultiLens MirFlickr Ambient Mini (67.7 MB) https://huggingface.co/datasets/Lensless/MultiLens-Mirflickr-Ambient-100
+    "multilens_mirflickr_ambient_mini": {
+        "size (GB)": 0.0677,
+        "huggingface_repo": "Lensless/MultiLens-Mirflickr-Ambient-100",
+        "psf": "psf.png",
+        "display_res": [600, 600],
+        "alignment": {"top_left": [118, 220], "height": 123},
+    },
+}
+
+
+def print_available_datasets():
+    print("Available datasets:")
+    for dataset in available_datasets:
+        print(
+            f"  - {dataset} ({available_datasets[dataset]['size (GB)']} GB) : https://huggingface.co/datasets/{available_datasets[dataset]['huggingface_repo']}"
+        )
+
+
+def get_dataset(dataset_name, split, cache_dir=None, **kwargs):
+    """
+    Get a dataset by name.
+
+    Parameters
+    ----------
+    dataset_name : str
+        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
+    -------
+    :py:class:`~torch.utils.data.Dataset`
+        Dataset object.
+    """
+    if dataset_name not in available_datasets:
+        print_str_available_dataset = "Available datasets are:"
+        for dataset in available_datasets:
+            print_str_available_dataset += f"\n  - {dataset} ({available_datasets[dataset]['size (GB)']} GB) : https://huggingface.co/datasets/{available_datasets[dataset]['huggingface_repo']}"
+        raise ValueError(
+            f"Dataset '{dataset_name}' not available.\n\n{print_str_available_dataset}"
+        )
+    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)
+
+
 class DualDataset(Dataset):
     """
     Abstract class for defining a dataset of paired lensed and lensless images.

diff --git a/lensless/utils/plot.py b/lensless/utils/plot.py
@@ -290,6 +290,38 @@ def plot_autocorr2d(vals, pad_mode="reflect", ax=None):
     return ax, autocorr
 
 
+def plot_autocorr_rgb(img, width=3, figsize=None):
+    """
+    Plot autocorrelation of each channel of an image.
+
+    Parameters
+    ----------
+    img : py:class:`~numpy.ndarray`
+        2-D image.
+    width : int, optional
+        Width of cross-section to plot. Default is 3dB.
+
+    """
+
+    assert len(img.shape) == 3, "Image must be 3D"
+    assert img.shape[2] == 3, "Image must have 3 color channels"
+
+    _, ax_auto = plt.subplots(ncols=3, nrows=2, num="Autocorrelations", figsize=figsize)
+
+    for i, c in enumerate(["r", "g", "b"]):
+        _, autocorr_c = plot_autocorr2d(img[:, :, i], ax=ax_auto[0][i])
+
+        ax, _ = plot_cross_section(
+            autocorr_c,
+            color=c,
+            ax=ax_auto[1][i],
+            plot_db_drop=width,
+        )
+        if i != 0:
+            ax.set_ylabel("")
+    return ax
+
+
 def compare_models(model_paths, max_epoch=None, linewidth=2, fontsize=18, metrics=None):
     """
     Plot train and test loss for multiple models, and print metrics for best epoch.

diff --git a/notebook/lenslesspicam_demo.ipynb b/notebook/lenslesspicam_demo.ipynb