Ilo

configs/ilo_single_file.yaml

@@ -0,0 +1,46 @@
+hydra:
+  job:
+    chdir: True    # change to output folder
+
+
+files:
+  original: data/celeba_mini/000019.jpg
+  #original: data/original/mnist_3.png
+  psf: data/psf/tape_rgb.png
+
+save: True
+
+simulation:
+  object_height: 0.3
+  # these distance parameters are typically fixed for a given PSF
+  scene2mask: 40e-2
+  mask2sensor: 4e-3
+  # see waveprop.devices
+  sensor: "rpi_hq"
+  snr_db: 20
+  # Downsampling for PSF
+  downsample: 8
+
+  # max val in simulated measured (quantized 8 bits)
+  max_val: 230
+
+  image_format: RGB
+
+  flatcam: False   # only supported if mask.type is "MURA" or "MLS"
+
+
+mask:
+  type: "tape"    # "MURA", "MLS", "FZA", "PhaseContour"
+
+  # Coded Aperture (MURA or MLS)
+  #flatcam_method: 'MLS'
+  n_bits: 8 # e.g. 8 for MLS, 99 for MURA
+
+  # Phase Contour
+  noise_period: [16, 16]
+  refractive_index: 1.2
+  phase_mask_iter: 10
+
+  # Fresnel Zone Aperture
+  radius: 0.32e-3
+

lensless/utils/align_face.py

@@ -0,0 +1,178 @@
+import os
+import numpy as np
+import PIL.Image
+import scipy.ndimage
+import dlib
+
+try:
+    import torch
+    import torchvision.transforms as tf
+    from torchvision.datasets.utils import download_and_extract_archive
+
+    torch_available = True
+except ImportError:
+    torch_available = False
+
+project_root = os.getcwd()
+marker = ".gitignore"
+while not os.path.exists(os.path.join(project_root, marker)):
+    project_root = os.path.dirname(project_root)
+    if project_root == os.path.dirname(project_root):
+        raise FileNotFoundError(
+            ".gitignore file not found. Are you sure you are in the 'Lensless' project?"
+        )
+model_dir = os.path.relpath(project_root, os.getcwd()) + os.sep + "data/models"
+if not os.path.isdir(model_dir):
+    os.makedirs(model_dir)
+if not os.path.exists(os.path.join(model_dir, "shape_predictor_68_face_landmarks.dat")):
+    msg = "Do you want to download the face landmark model (61.1 Mo)?"
+    valid = input("%s (Y/n) " % msg).lower() != "n"
+    if valid:
+        url = "http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2"
+        filename = "shape_predictor_68_face_landmarks.dat.bz2"
+        download_and_extract_archive(url, model_dir, filename=filename, remove_finished=True)
+
+predictor = dlib.shape_predictor(os.path.join(model_dir, "shape_predictor_68_face_landmarks.dat"))
+
+
+def get_landmark(filepath):
+    """get landmark with dlib
+    :return: np.array shape=(68, 2)
+    """
+    detector = dlib.get_frontal_face_detector()
+
+    img = dlib.load_rgb_image(filepath)
+    dets = detector(img, 1)
+
+    print("Number of faces detected: {}".format(len(dets)))
+    for k, d in enumerate(dets):
+        print(
+            "Detection {}: Left: {} Top: {} Right: {} Bottom: {}".format(
+                k, d.left(), d.top(), d.right(), d.bottom()
+            )
+        )
+        # Get the landmarks/parts for the face in box d.
+        shape = predictor(img, d)
+        print("Part 0: {}, Part 1: {} ...".format(shape.part(0), shape.part(1)))
+
+    t = list(shape.parts())
+    a = []
+    for tt in t:
+        a.append([tt.x, tt.y])
+    lm = np.array(a)
+    # lm is a shape=(68,2) np.array
+    return lm
+
+
+def align_face(filepath):
+    """
+    :param filepath: str
+    :return: PIL Image
+    """
+
+    lm = get_landmark(filepath)
+
+    # lm_chin = lm[0:17]  # left-right
+    # lm_eyebrow_left = lm[17:22]  # left-right
+    # lm_eyebrow_right = lm[22:27]  # left-right
+    # lm_nose = lm[27:31]  # top-down
+    # lm_nostrils = lm[31:36]  # top-down
+    lm_eye_left = lm[36:42]  # left-clockwise
+    lm_eye_right = lm[42:48]  # left-clockwise
+    lm_mouth_outer = lm[48:60]  # left-clockwise
+    # lm_mouth_inner = lm[60:68]  # left-clockwise
+
+    # Calculate auxiliary vectors.
+    eye_left = np.mean(lm_eye_left, axis=0)
+    eye_right = np.mean(lm_eye_right, axis=0)
+    eye_avg = (eye_left + eye_right) * 0.5
+    eye_to_eye = eye_right - eye_left
+    mouth_left = lm_mouth_outer[0]
+    mouth_right = lm_mouth_outer[6]
+    mouth_avg = (mouth_left + mouth_right) * 0.5
+    eye_to_mouth = mouth_avg - eye_avg
+
+    # Choose oriented crop rectangle.
+    x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
+    x /= np.hypot(*x)
+    x *= max(np.hypot(*eye_to_eye) * 2.0, np.hypot(*eye_to_mouth) * 1.8)
+    y = np.flipud(x) * [-1, 1]
+    c = eye_avg + eye_to_mouth * 0.1
+    quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
+    qsize = np.hypot(*x) * 2
+
+    # read image
+    img = PIL.Image.open(filepath)
+
+    output_size = 1024
+    transform_size = 4096
+    enable_padding = True
+
+    # Shrink.
+    shrink = int(np.floor(qsize / output_size * 0.5))
+    if shrink > 1:
+        rsize = (
+            int(np.rint(float(img.size[0]) / shrink)),
+            int(np.rint(float(img.size[1]) / shrink)),
+        )
+        img = img.resize(rsize, PIL.Image.ANTIALIAS)
+        quad /= shrink
+        qsize /= shrink
+
+    # Crop.
+    border = max(int(np.rint(qsize * 0.1)), 3)
+    crop = (
+        int(np.floor(min(quad[:, 0]))),
+        int(np.floor(min(quad[:, 1]))),
+        int(np.ceil(max(quad[:, 0]))),
+        int(np.ceil(max(quad[:, 1]))),
+    )
+    crop = (
+        max(crop[0] - border, 0),
+        max(crop[1] - border, 0),
+        min(crop[2] + border, img.size[0]),
+        min(crop[3] + border, img.size[1]),
+    )
+    if crop[2] - crop[0] < img.size[0] or crop[3] - crop[1] < img.size[1]:
+        img = img.crop(crop)
+        quad -= crop[0:2]
+
+    # Pad.
+    pad = (
+        int(np.floor(min(quad[:, 0]))),
+        int(np.floor(min(quad[:, 1]))),
+        int(np.ceil(max(quad[:, 0]))),
+        int(np.ceil(max(quad[:, 1]))),
+    )
+    pad = (
+        max(-pad[0] + border, 0),
+        max(-pad[1] + border, 0),
+        max(pad[2] - img.size[0] + border, 0),
+        max(pad[3] - img.size[1] + border, 0),
+    )
+    if enable_padding and max(pad) > border - 4:
+        pad = np.maximum(pad, int(np.rint(qsize * 0.3)))
+        img = np.pad(np.float32(img), ((pad[1], pad[3]), (pad[0], pad[2]), (0, 0)), "reflect")
+        h, w, _ = img.shape
+        y, x, _ = np.ogrid[:h, :w, :1]
+        mask = np.maximum(
+            1.0 - np.minimum(np.float32(x) / pad[0], np.float32(w - 1 - x) / pad[2]),
+            1.0 - np.minimum(np.float32(y) / pad[1], np.float32(h - 1 - y) / pad[3]),
+        )
+        blur = qsize * 0.02
+        img += (scipy.ndimage.gaussian_filter(img, [blur, blur, 0]) - img) * np.clip(
+            mask * 3.0 + 1.0, 0.0, 1.0
+        )
+        img += (np.median(img, axis=(0, 1)) - img) * np.clip(mask, 0.0, 1.0)
+        img = PIL.Image.fromarray(np.uint8(np.clip(np.rint(img), 0, 255)), "RGB")
+        quad += pad[:2]
+
+    # Transform.
+    img = img.transform(
+        (transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR
+    )
+    if output_size < transform_size:
+        img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS)
+
+    # Save aligned image.
+    return img

scripts/sim/ilo_single_file.py

@@ -0,0 +1,167 @@
+"""
+
+Simulate a mask, use face alignment on a face image and simulate a measurement with the mask on the image.
+
+Procedure is as follows:
+
+1) Simulate the mask.
+2) Align the face.
+3) Simulate a measurement with the mask and specified physical parameters.
+
+Example usage:
+
+Simulate FlatCam with separable simulation (https://arxiv.org/abs/1509.00116, Eq 7):
+```
+python scripts/sim/ilo_single_file.py mask.type=MLS simulation.flatcam=True recon.algo=tikhonov
+```
+
+Simulate FlatCam with PSF simulation:
+```
+python scripts/sim/ilo_single_file.py mask.type=MLS simulation.flatcam=False
+```
+
+Simulate Fresnel Zone Aperture camera with PSF simulation (https://www.nature.com/articles/s41377-020-0289-9):
+```
+python scripts/sim/ilo_single_file.py mask.type=FZA
+```
+
+Simulate PhaseContour camera with PSF simulation (https://ieeexplore.ieee.org/document/9076617):
+```
+python scripts/sim/ilo_single_file.py mask.type=PhaseContour
+```
+
+"""
+
+import hydra
+import warnings
+from hydra.utils import to_absolute_path
+from lensless.utils.io import load_psf  # , save_image
+from lensless.utils.image import rgb2gray, rgb2bayer  # , align_face
+from lensless.utils.align_face import align_face
+import numpy as np
+import matplotlib.pyplot as plt
+from lensless.utils.plot import plot_image
+
+# from lensless.eval.metric import mse, psnr, ssim, lpips
+from waveprop.simulation import FarFieldSimulator
+import os
+from lensless.hardware.mask import CodedAperture, PhaseContour, FresnelZoneAperture
+
+
+@hydra.main(version_base=None, config_path="../../configs", config_name="ilo_single_file")
+def simulate(config):
+
+    fp = to_absolute_path(config.files.original)
+    assert os.path.exists(fp), f"File {fp} does not exist."
+
+    # simulation parameters
+    object_height = config.simulation.object_height
+    scene2mask = config.simulation.scene2mask
+    mask2sensor = config.simulation.mask2sensor
+    sensor = config.simulation.sensor
+    snr_db = config.simulation.snr_db
+    downsample = config.simulation.downsample
+    max_val = config.simulation.max_val
+
+    image_format = config.simulation.image_format.lower()
+    grayscale = False
+    if image_format == "grayscale":
+        grayscale = True
+
+    # 1) simulate mask
+    mask_type = config.mask.type
+    if mask_type.upper() in ["MURA", "MLS"]:
+        mask = CodedAperture.from_sensor(
+            sensor_name=sensor,
+            downsample=downsample,
+            method=mask_type,
+            distance_sensor=mask2sensor,
+            **config.mask,
+        )
+        psf = mask.psf / np.linalg.norm(mask.psf.ravel())
+    elif mask_type.upper() == "FZA":
+        mask = FresnelZoneAperture.from_sensor(
+            sensor_name=sensor,
+            downsample=downsample,
+            distance_sensor=mask2sensor,
+            **config.mask,
+        )
+        psf = mask.psf / np.linalg.norm(mask.psf.ravel())
+    elif mask_type == "PhaseContour":
+        mask = PhaseContour.from_sensor(
+            sensor_name=sensor,
+            downsample=downsample,
+            distance_sensor=mask2sensor,
+            **config.mask,
+        )
+        psf = mask.psf / np.linalg.norm(mask.psf.ravel())
+    else:
+        psf_fp = to_absolute_path(config.files.psf)
+        assert os.path.exists(psf_fp), f"PSF {psf_fp} does not exist."
+        psf = load_psf(psf_fp, verbose=True, downsample=downsample)
+        psf = psf.squeeze()
+
+    if grayscale and psf.shape[-1] == 3:
+        psf = rgb2gray(psf)
+    if downsample > 1:
+        print(f"Downsampled to {psf.shape}.")
+
+    # 2) simulate measurement
+    # image = load_image(fp, verbose=True)
+    image = np.array(align_face(fp)) / 255
+    if grayscale and len(image.shape) == 3:
+        image = rgb2gray(image)
+
+    flatcam_sim = config.simulation.flatcam
+    if flatcam_sim and mask_type.upper() not in ["MURA", "MLS"]:
+        warnings.warn(
+            "Flatcam simulation only supported for MURA and MLS masks. Using far field simulation with PSF."
+        )
+        flatcam_sim = False
+
+    # use far field simulator to get correct object plane sizing
+    simulator = FarFieldSimulator(
+        psf=psf,  # only support one depth plane
+        object_height=object_height,
+        scene2mask=scene2mask,
+        mask2sensor=mask2sensor,
+        sensor=sensor,
+        snr_db=snr_db,
+        max_val=max_val,
+    )
+    image_plane, object_plane = simulator.propagate(image, return_object_plane=True)
+
+    if image_format == "grayscale":
+        image_plane = rgb2gray(image_plane)
+        object_plane = rgb2gray(object_plane)
+    elif "bayer" in image_format:
+        image_plane = rgb2bayer(image_plane, pattern=image_format[-4:])
+        object_plane = rgb2bayer(object_plane, pattern=image_format[-4:])
+    else:
+        # make sure image is RGB
+        assert image_plane.shape[-1] == 3, "Image plane must be RGB"
+        assert object_plane.shape[-1] == 3, "Object plane must be RGB"
+
+    if flatcam_sim:
+        # apply flatcam simulation to object plane
+        image_plane = mask.simulate(object_plane, snr_db=snr_db)
+
+    # -- plot
+    fig, ax = plt.subplots(ncols=3, nrows=1, figsize=(15, 5))
+    plot_image(object_plane, ax=ax[0])
+    ax[0].set_title("Object plane")
+    plot_image(psf, ax=ax[1], gamma=3.5)
+    ax[1].set_title("PSF")
+    plot_image(image_plane, ax=ax[2])
+    ax[2].set_title("Raw data")
+    plt.savefig("result.png")
+    fig.tight_layout()
+
+    for a in ax:
+        a.set_axis_off()
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    simulate()