fft2d 0.1.0

Fourier transform for 2D data such as images
Documentation 
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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

// SPDX-License-Identifier: MPL-2.0

//! Fourier transform for 2D data such as images.

use rustfft::FftDirection;
use rustfft::{num_complex::Complex, FftPlanner};

/// Compute the 2D Fourier transform of an image buffer.
///
/// The image buffer is considered to be stored in row major order.
/// After the 2D FFT has been applied, the buffer contains the transposed
/// of the Fourier transform since one transposition is needed to process
/// the columns of the image buffer.
///
/// The transformation is not normalized.
/// To normalize the output, you should multiply it by 1 / sqrt( width * height ).
/// If the transformed buffer is intended to be processed
/// and then converted back into an image with an inverse Fourier transform,
/// it is more efficient to multiply at the end by 1 / (width * height).
///
/// Remark: an allocation the size of the image buffer is performed for the transposition,
/// as well as scratch buffers while performing the rows and columns FFTs.
pub fn fft_2d(width: usize, height: usize, img_buffer: &mut [Complex<f64>]) {
    fft_2d_with_direction(width, height, img_buffer, FftDirection::Forward)
}

/// Compute the inverse 2D Fourier transform to get back an image buffer.
///
/// After the inverse 2D FFT has been applied, the image buffer contains the transposed
/// of the inverse Fourier transform since one transposition is needed to process
/// the columns of the buffer.
///
/// The transformation is not normalized.
/// To normalize the output, you should multiply it by 1 / sqrt( width * height ).
/// If this is used as a pair of FFT followed by inverse FFT,
/// is is more efficient to normalize only once by 1 / (width * height) at the end.
///
/// Remark: an allocation the size of the image buffer is performed for the transposition,
/// as well as scratch buffers while performing the rows and columns FFTs.
pub fn ifft_2d(width: usize, height: usize, img_buffer: &mut [Complex<f64>]) {
    fft_2d_with_direction(width, height, img_buffer, FftDirection::Inverse)
}

/// Compute the 2D Fourier transform or inverse transform of an image buffer.
///
/// The image buffer is considered to be stored in row major order.
/// After the 2D FFT has been applied, the buffer contains the transposed
/// of the Fourier transform since one transposition is needed to process
/// the columns of the image buffer.
///
/// The transformation is not normalized.
/// To normalize the output, you should multiply it by 1 / sqrt( width * height ).
/// If this is used as a pair of FFT followed by inverse FFT,
/// is is more efficient to normalize only once by 1 / (width * height) at the end.
///
/// Remark: an allocation the size of the image buffer is performed for the transposition,
/// as well as a scratch buffer while performing the rows and columns FFTs.
fn fft_2d_with_direction(
    width: usize,
    height: usize,
    img_buffer: &mut [Complex<f64>],
    direction: FftDirection,
) {
    // Compute the FFT of each row of the image.
    let mut planner = FftPlanner::new();
    let fft_width = planner.plan_fft(width, direction);
    let mut scratch = vec![Complex::default(); fft_width.get_inplace_scratch_len()];
    for row_buffer in img_buffer.chunks_exact_mut(width) {
        fft_width.process_with_scratch(row_buffer, &mut scratch);
    }

    // Transpose the image to be able to compute the FFT on the other dimension.
    let mut transposed = transpose(width, height, img_buffer);
    let fft_height = planner.plan_fft(height, direction);
    scratch.resize(fft_height.get_outofplace_scratch_len(), Complex::default());
    for (tr_buf, col_buf) in transposed
        .chunks_exact_mut(height)
        .zip(img_buffer.chunks_exact_mut(height))
    {
        fft_height.process_outofplace_with_scratch(tr_buf, col_buf, &mut scratch);
    }
}

fn transpose<T: Copy + Default>(width: usize, height: usize, matrix: &[T]) -> Vec<T> {
    let mut ind = 0;
    let mut ind_tr;
    let mut transposed = vec![T::default(); matrix.len()];
    for row in 0..height {
        ind_tr = row;
        for _ in 0..width {
            transposed[ind_tr] = matrix[ind];
            ind += 1;
            ind_tr += height;
        }
    }
    transposed
}

/// Inverse operation of the quadrants shift performed by fftshift.
///
/// It is different than fftshift if one dimension has an odd length.
pub fn ifftshift<T: Copy + Default>(width: usize, height: usize, matrix: &[T]) -> Vec<T> {
    // TODO: do actual code instead of relying on fftshift.
    let is_even = |length| length % 2 == 0;
    assert!(is_even(width), "Need a dedicated implementation");
    assert!(is_even(height), "Need a dedicated implementation");
    fftshift(width, height, matrix)
}

/// Shift the 4 quadrants of a Fourier transform to have all the low frequencies
/// at the center of the image.
pub fn fftshift<T: Copy + Default>(width: usize, height: usize, matrix: &[T]) -> Vec<T> {
    let mut shifted = vec![T::default(); matrix.len()];
    let half_width = width / 2;
    let half_height = height / 2;
    // Shift top and bottom quadrants.
    for row in 0..half_height {
        // top
        let mrow_start = row * width;
        let m_row = &matrix[mrow_start..mrow_start + width];
        // bottom
        let srow_start = mrow_start + (height - half_height) * width;
        let s_row = &mut shifted[srow_start..srow_start + width];
        // swap left and right
        s_row[width - half_width..width].copy_from_slice(&m_row[0..half_width]);
        s_row[0..width - half_width].copy_from_slice(&m_row[half_width..width]);
    }
    // Shift bottom and top quadrants.
    for row in half_height..height {
        // bottom
        let mrow_start = row * width;
        let m_row = &matrix[mrow_start..mrow_start + width];
        // top
        let srow_start = (row - half_height) * width;
        let s_row = &mut shifted[srow_start..srow_start + width];
        // swap left and right
        s_row[width - half_width..width].copy_from_slice(&m_row[0..half_width]);
        s_row[0..width - half_width].copy_from_slice(&m_row[half_width..width]);
    }
    shifted
}

// Sine and Cosine transforms ##################################################

#[cfg(feature = "rustdct")]
/// Cosine and Sine transforms.
pub mod dcst {

    use super::transpose;
    use rustdct::DctPlanner;

    /// Compute the 2D cosine transform of an image buffer.
    ///
    /// The image buffer is considered to be stored in row major order.
    /// After the 2D DCT has been applied, the buffer contains the transposed
    /// of the cosine transform since one transposition is needed to process
    /// the columns of the image buffer.
    ///
    /// The transformation is not normalized.
    /// To normalize the output, you should multiply it by 2 / sqrt( width * height ).
    /// If this is used as a pair of DCT followed by inverse DCT,
    /// is is more efficient to normalize only once at the end.
    ///
    /// Remark: an allocation the size of the image buffer is performed for the transposition,
    /// as well as a scratch buffer while performing the rows and columns transforms.
    pub fn dct_2d(width: usize, height: usize, img_buffer: &mut [f64]) {
        // Compute the FFT of each row of the image.
        let mut planner = DctPlanner::new();
        let dct_width = planner.plan_dct2(width);
        let mut scratch = vec![0.0; dct_width.get_scratch_len()];
        for row_buffer in img_buffer.chunks_exact_mut(width) {
            dct_width.process_dct2_with_scratch(row_buffer, &mut scratch);
        }

        // Transpose the image to be able to compute the FFT on the other dimension.
        let mut transposed = transpose(width, height, img_buffer);
        let dct_height = planner.plan_dct2(height);
        scratch.resize(dct_height.get_scratch_len(), 0.0);
        for column_buffer in transposed.chunks_exact_mut(height) {
            dct_height.process_dct2_with_scratch(column_buffer, &mut scratch);
        }
        img_buffer.copy_from_slice(&transposed);
    }

    /// Compute the inverse 2D cosine transform of an image buffer.
    ///
    /// The image buffer is considered to be stored in row major order.
    /// After the 2D iDCT has been applied, the buffer contains the transposed
    /// of the cosine transform since one transposition is needed to process
    /// the columns of the image buffer.
    ///
    /// The transformation is not normalized.
    /// To normalize the output, you should multiply it by 2 / sqrt( width * height ).
    /// If this is used as a pair of DCT followed by inverse DCT,
    /// is is more efficient to normalize only once at the end.
    ///
    /// Remark: an allocation the size of the image buffer is performed for the transposition,
    /// as well as a scratch buffer while performing the rows and columns transforms.
    pub fn idct_2d(width: usize, height: usize, img_buffer: &mut [f64]) {
        // Compute the FFT of each row of the image.
        let mut planner = DctPlanner::new();
        let dct_width = planner.plan_dct3(width);
        let mut scratch = vec![0.0; dct_width.get_scratch_len()];
        for row_buffer in img_buffer.chunks_exact_mut(width) {
            dct_width.process_dct3_with_scratch(row_buffer, &mut scratch);
        }

        // Transpose the image to be able to compute the FFT on the other dimension.
        let mut transposed = transpose(width, height, img_buffer);
        let dct_height = planner.plan_dct3(height);
        scratch.resize(dct_height.get_scratch_len(), 0.0);
        for column_buffer in transposed.chunks_exact_mut(height) {
            dct_height.process_dct3_with_scratch(column_buffer, &mut scratch);
        }
        img_buffer.copy_from_slice(&transposed);
    }
}