.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/transforms-equivariance/demo_transforms.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note New to DeepInverse? Get started with the basics with the :ref:`5 minute quickstart tutorial `.. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_transforms-equivariance_demo_transforms.py: Image transforms for equivariance & augmentations ================================================= We demonstrate the use of our ``deepinv.transform`` module for use in solving imaging problems. These can be used for: 1. Data augmentation (similar to ``torchvision.transforms``) 2. Building equivariant denoisers (:class:`deepinv.models.EquivariantDenoiser`) for robust denoising (e.g from :footcite:t:`terris2024equivariant`) 3. Self-supervised learning using Equivariant Imaging from :footcite:t:`chen2021equivariant`. See :ref:`sphx_glr_auto_examples_self-supervised-learning_demo_ei_transforms.py`, :ref:`sphx_glr_auto_examples_self-supervised-learning_demo_equivariant_imaging.py` for thorough examples. See :ref:`docs ` for full list of implemented transforms. 1. Data augmentation -------------------- We can use ``deepinv`` transforms in the same way as ``torchvision`` transforms, and chain them together for data augmentation. Our transforms are customisable and offer some group-theoretic properties. We demonstrate a random roto-scale combined with a random masking, and a constrained pixel-shift with a random color jitter. Note that all our transforms can easily be inverted using the method ``transform.inverse()``. First, load a sample image. .. GENERATED FROM PYTHON SOURCE LINES 33-65 .. code-block:: Python import deepinv as dinv import torch from torchvision.transforms import Compose, ColorJitter, RandomErasing, Resize x = dinv.utils.load_example("celeba_example.jpg") # Random roto-scale with random masking transform = Compose( [ dinv.transform.Rotate() * dinv.transform.Scale(), RandomErasing(), ] ) # Constrained pixel-shift with a random color jitter transform2 = Compose( [ dinv.transform.Shift(shift_max=0.2), ColorJitter(hue=0.5), ] ) # Random diffeomorphism transform3 = dinv.transform.CPABDiffeomorphism() dinv.utils.plot( [x, transform(x), transform2(x), transform3(x)], titles=["Orig", "Transform 1", "Transform 2", "Transform 3"], ) .. image-sg:: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_001.png :alt: Orig, Transform 1, Transform 2, Transform 3 :srcset: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none /local/jtachell/deepinv/deepinv/deepinv/transform/rotate.py:49: UserWarning: The default interpolation mode will be changed to bilinear interpolation in the near future. Please specify the interpolation mode explicitly if you plan to keep using nearest interpolation. warn( .. GENERATED FROM PYTHON SOURCE LINES 66-69 By letting ``n_trans`` be equal to the full group size, all transforms are recovered: .. GENERATED FROM PYTHON SOURCE LINES 69-77 .. code-block:: Python reflect = dinv.transform.Reflect(dim=[-2, -1], n_trans=4) rotate = dinv.transform.Rotate(multiples=90, positive=True, n_trans=4) dinv.utils.plot( [reflect(x), rotate(x)], titles=["Full 2D reflect group", "Full rotate group"] ) .. image-sg:: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_002.png :alt: Full 2D reflect group, Full rotate group :srcset: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 78-86 2. Equivariant denoiser or plug-and-play ---------------------------------------- Suppose we want to make a denoiser equivariant to the rotoreflect group, taken as the group product of the 90 degree rotations (order 4) and 1D reflects (order 2). We can do this with our transform arithmetic (note this results in the full dihedral group :math:`\text{Dih}_4` of order 8): .. GENERATED FROM PYTHON SOURCE LINES 86-90 .. code-block:: Python transform = rotate * dinv.transform.Reflect(dim=[-1], n_trans=2) .. GENERATED FROM PYTHON SOURCE LINES 91-95 Let's simulate some Gaussian noise and turn a simple (median filter) denoiser into an equivariant denoiser (:class:`deepinv.models.EquivariantDenoiser`): .. GENERATED FROM PYTHON SOURCE LINES 95-128 .. code-block:: Python sigma = 0.1 physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=sigma)) x = Resize(128)(x) # Put the image in an unusual orientation to show the benefits of equivariance x = torch.rot90(x, k=1, dims=[-2, -1]) y = physics(x) model = dinv.models.RAM(pretrained=True) model_eq = dinv.models.EquivariantDenoiser(model, transform=transform) with torch.no_grad(): x_hat = model(y, sigma=sigma) x_eq = model_eq(y, sigma=sigma) psnr_fn = dinv.metric.PSNR() psnr = psnr_fn(x_hat, x).item() psnr_eq = psnr_fn(x_eq, x).item() psnr_y = psnr_fn(y, x).item() dinv.utils.plot( [y, x_hat, x_eq, x], ["Measurements", "Regular Denoiser", "Equivariant Denoiser", "Ground truth"], subtitles=[ f"PSNR={psnr_y:.1f}dB", f"PSNR={psnr:.1f}dB", f"PSNR={psnr_eq:.1f}dB", "", ], fontsize=10, ) .. image-sg:: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_003.png :alt: Measurements, Regular Denoiser, Equivariant Denoiser, Ground truth :srcset: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 129-133 What's going on under the hood? We use the ``transform.symmetrize`` method to symmetrize the function :math:`D` with respect to a projective transform (with a Monte Carlo approach of ``n_trans=2`` transforms per call): .. GENERATED FROM PYTHON SOURCE LINES 133-146 .. code-block:: Python # Example non-equivariant function D = lambda x: x[..., [0]] * x # Example non-linear transform with n=2 t = dinv.transform.projective.PanTiltRotate(n_trans=2, theta_max=10, theta_z_max=0) # Symmetrize function with respect to transform D_s = t.symmetrize(D, average=True) dinv.utils.plot( [x, D(x), D_s(x)], titles=["Orig", "$D(x)$", "$\\sum_i T_g^{-1}D(T_g x)$"] ) .. image-sg:: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_004.png :alt: Orig, $D(x)$, $\sum_i T_g^{-1}D(T_g x)$ :srcset: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 147-154 Reconstructors can also be made equivariant in a similar way using :class:`deepinv.models.EquivariantReconstructor`. This amounts to averaging the base reconstructor :math:`\tilde{R}` over the transformations to get an equivariant reconstructor .. math:: R(y, A) = \frac{1}{|\mathcal{G}|}\sum_{g\in \mathcal{G}} T_g \tilde{R}(y, A T_g) which is computed using a Monte Carlo sampling where a random subset of transformations is used, typically a single one at training time and the full set at evaluation time. .. GENERATED FROM PYTHON SOURCE LINES 154-190 .. code-block:: Python sigma = 0.1 physics = dinv.physics.Denoising(noise_model=dinv.physics.GaussianNoise(sigma=sigma)) y = physics(x) rotate = dinv.transform.Rotate(multiples=90, positive=True, n_trans=4) transform = rotate * dinv.transform.Reflect(dim=[-1], n_trans=2) model = dinv.models.RAM(pretrained=True) model_eq = dinv.models.EquivariantReconstructor(model, transform=transform) with torch.no_grad(): x_hat = model(y, physics=physics) x_eq = model_eq(y, physics=physics) psnr = psnr_fn(x_hat, x).item() psnr_eq = psnr_fn(x_eq, x).item() psnr_y = psnr_fn(y, x).item() dinv.utils.plot( [y, x_hat, x_eq, x], [ "Measurements", "Regular Reconstructor", "Equivariant Reconstructor", "Ground truth", ], subtitles=[ f"PSNR={psnr_y:.1f}dB", f"PSNR={psnr:.1f}dB", f"PSNR={psnr_eq:.1f}dB", "", ], fontsize=9, ) .. image-sg:: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_005.png :alt: Measurements, Regular Reconstructor, Equivariant Reconstructor, Ground truth :srcset: /auto_examples/transforms-equivariance/images/sphx_glr_demo_transforms_005.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 191-202 3. Equivariant imaging ---------------------- We can also use our transforms to create the self-supervised equivariant imaging loss. See :ref:`sphx_glr_auto_examples_self-supervised-learning_demo_ei_transforms.py`, :ref:`sphx_glr_auto_examples_self-supervised-learning_demo_equivariant_imaging.py` for examples of self-supervised learning for MRI and inpainting. For example, the EI loss can easily be defined using any combination of transforms: .. GENERATED FROM PYTHON SOURCE LINES 202-207 .. code-block:: Python loss = dinv.loss.EILoss( transform=dinv.transform.projective.Affine() | dinv.transform.projective.Euclidean() ) .. GENERATED FROM PYTHON SOURCE LINES 208-211 :References: .. footbibliography:: .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 13.121 seconds) .. _sphx_glr_download_auto_examples_transforms-equivariance_demo_transforms.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: demo_transforms.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: demo_transforms.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: demo_transforms.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_