phase-retrieval/phase_retrieval.py at master · tuelwer/phase-retrieval · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import numpy as np

def fienup_phase_retrieval(mag, mask=None, beta=0.8,
                           steps=200, mode='hybrid', verbose=True):
    """
    Implementation of Fienup's phase-retrieval methods. This function
    implements the input-output, the output-output and the hybrid method.

    Note: Mode 'output-output' and beta=1 results in
    the Gerchberg-Saxton algorithm.

    Parameters:
        mag: Measured magnitudes of Fourier transform
        mask: Binary array indicating where the image should be
              if padding is known
        beta: Positive step size
        steps: Number of iterations
        mode: Which algorithm to use
              (can be 'input-output', 'output-output' or 'hybrid')
        verbose: If True, progress is shown

    Returns:
        x: Reconstructed image

    Author: Tobias Uelwer
    Date: 30.12.2018

    References:
    [1] E. Osherovich, Numerical methods for phase retrieval, 2012,
        https://arxiv.org/abs/1203.4756
    [2] J. R. Fienup, Phase retrieval algorithms: a comparison, 1982,
        https://www.osapublishing.org/ao/abstract.cfm?uri=ao-21-15-2758
    [3] https://github.com/cwg45/Image-Reconstruction
    """

    assert beta > 0, 'step size must be a positive number'
    assert steps > 0, 'steps must be a positive number'
    assert mode == 'input-output' or mode == 'output-output'\
        or mode == 'hybrid',\
    'mode must be \'input-output\', \'output-output\' or \'hybrid\''

    if mask is None:
        mask = np.ones(mag.shape)

    assert mag.shape == mask.shape, 'mask and mag must have same shape'

    # sample random phase and initialize image x
    y_hat = mag*np.exp(1j*2*np.pi*np.random.rand(*mag.shape))
    x = np.zeros(mag.shape)

    # previous iterate
    x_p = None

    # main loop
    for i in range(1, steps+1):
        # show progress
        if i % 100 == 0 and verbose:
            print("step", i, "of", steps)

        # inverse fourier transform
        y = np.real(np.fft.ifft2(y_hat))

        # previous iterate
        if x_p is None:
            x_p = y
        else:
            x_p = x

        # updates for elements that satisfy object domain constraints
        if mode == "output-output" or mode == "hybrid":
            x = y

        # find elements that violate object domain constraints
        # or are not masked
        indices = np.logical_or(np.logical_and(y<0, mask),
                                np.logical_not(mask))

        # updates for elements that violate object domain constraints
        if mode == "hybrid" or mode == "input-output":
            x[indices] = x_p[indices]-beta*y[indices]
        elif mode == "output-output":
            x[indices] = y[indices]-beta*y[indices]

        # fourier transform
        x_hat = np.fft.fft2(x)

        # satisfy fourier domain constraints
        # (replace magnitude with input magnitude)
        y_hat = mag*np.exp(1j*np.angle(x_hat))
    return x