|
| 1 | +""" |
| 2 | +Explore fourier transform of images |
| 3 | +=================================== |
| 4 | +Example showing how to use a `RectangleSelector` to interactively reconstruct |
| 5 | +an image using portions of it fourier transform |
| 6 | +""" |
| 7 | + |
| 8 | +# test_example = false |
| 9 | +# sphinx_gallery_pygfx_docs = 'screenshot' |
| 10 | + |
| 11 | +import numpy as np |
| 12 | +import fastplotlib as fpl |
| 13 | +import imageio.v3 as iio |
| 14 | + |
| 15 | +image = iio.imread("imageio:camera.png") |
| 16 | + |
| 17 | +# compute discrete fourier transform of image |
| 18 | +img_fft = np.fft.fftshift(np.fft.fft2(image)) |
| 19 | + |
| 20 | +# image to just visualize absolute magnitudes |
| 21 | +log_abs_img_fft = np.log(np.abs(img_fft + 1)) |
| 22 | + |
| 23 | +# placeholders for displaying fft and inverse fft of selections |
| 24 | +zeros = np.zeros(image.shape) |
| 25 | + |
| 26 | +# create an ImageWidget to display all the images |
| 27 | +iw = fpl.ImageWidget( |
| 28 | + data=[image, log_abs_img_fft, zeros, zeros, zeros, zeros], |
| 29 | + names=["image", "DFT", "selected", "FFT of selected", "not-selected", "IFFT of not-selected"], |
| 30 | + figure_shape=(3, 2), # so we can see image and fft side by side |
| 31 | + figure_kwargs={"size": (700, 1024)}, |
| 32 | + histogram_widget=False, |
| 33 | +) |
| 34 | + |
| 35 | +# we don't need the toolbars here, unclutter the figure |
| 36 | +for subplot in iw.figure: |
| 37 | + subplot.toolbar = False |
| 38 | + |
| 39 | +# viridis cmap for the fft images |
| 40 | +iw.cmap = "viridis" |
| 41 | + |
| 42 | +# gray for the non-fft images |
| 43 | +iw.managed_graphics[0].cmap = "gray" |
| 44 | +iw.managed_graphics[3].cmap = "gray" |
| 45 | +iw.managed_graphics[-1].cmap = "gray" |
| 46 | + |
| 47 | +# set contrast limits based on the full DFT for the DFT-selection images |
| 48 | +iw.figure["selected"].graphics[0].vmin, iw.figure["selected"].graphics[0].vmax = log_abs_img_fft.min(), log_abs_img_fft.max() |
| 49 | +iw.figure["not-selected"].graphics[0].vmin, iw.figure["not-selected"].graphics[0].vmax = log_abs_img_fft.min(), log_abs_img_fft.max() |
| 50 | + |
| 51 | +iw.show() |
| 52 | + |
| 53 | +# create a rectangle selector |
| 54 | +rs = iw.managed_graphics[1].add_rectangle_selector(edge_color="w", edge_thickness=2.0) |
| 55 | + |
| 56 | + |
| 57 | +@rs.add_event_handler("selection") |
| 58 | +def update_images(ev): |
| 59 | + """ |
| 60 | + Updates the images when the selection changes |
| 61 | + """ |
| 62 | + |
| 63 | + # get the bbox of the selection |
| 64 | + row_ixs, col_ixs = ev.get_selected_indices() |
| 65 | + row_slice = slice(row_ixs[0], row_ixs[-1] + 1) |
| 66 | + col_slice = slice(col_ixs[0], col_ixs[-1] + 1) |
| 67 | + |
| 68 | + # fft of the selection |
| 69 | + selected_fft = np.zeros(image.shape, dtype=np.complex64) |
| 70 | + selected_fft[row_slice, col_slice] = img_fft[row_slice, col_slice] |
| 71 | + |
| 72 | + # update image graphic with the current fft selection |
| 73 | + iw.managed_graphics[2].data = np.log(np.abs(selected_fft + 1)) |
| 74 | + |
| 75 | + # inverse fft to reconstruct image using only the selection |
| 76 | + iw.managed_graphics[3].data = np.fft.ifft2(np.fft.fftshift(selected_fft)) |
| 77 | + iw.managed_graphics[3].reset_vmin_vmax() |
| 78 | + |
| 79 | + # fft of the region outside the selection |
| 80 | + unselected_fft = img_fft.copy() |
| 81 | + unselected_fft[row_slice, col_slice] = 0 |
| 82 | + |
| 83 | + # update image graphic with unselected fft area |
| 84 | + iw.managed_graphics[4].data = np.log(np.abs(unselected_fft + 1)) |
| 85 | + |
| 86 | + # inverse fft to reconstruct image using only the unselected part of the fft |
| 87 | + iw.managed_graphics[5].data = np.fft.ifft2(np.fft.fftshift(unselected_fft)) |
| 88 | + iw.managed_graphics[5].reset_vmin_vmax() |
| 89 | + |
| 90 | + |
| 91 | +# set initial selection to the center |
| 92 | +rs.selection = (225, 285, 225, 285) |
| 93 | + |
| 94 | + |
| 95 | +figure = iw.figure |
| 96 | + |
| 97 | +# NOTE: `if __name__ == "__main__"` is NOT how to use fastplotlib interactively |
| 98 | +# please see our docs for using fastplotlib interactively in ipython and jupyter |
| 99 | +if __name__ == "__main__": |
| 100 | + print(__doc__) |
| 101 | + fpl.loop.run() |
0 commit comments