.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "auto_examples/self-supervised-learning/demo_ei_transforms.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

    .. note::
        :class: sphx-glr-download-link-note

        :ref:`Go to the end <sphx_glr_download_auto_examples_self-supervised-learning_demo_ei_transforms.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_auto_examples_self-supervised-learning_demo_ei_transforms.py:


Image transformations for Equivariant Imaging
=============================================

This example demonstrates various geometric image transformations
implemented in ``deepinv`` that can be used in Equivariant Imaging (EI)
for self-supervised learning:

-  Shift: integer pixel 2D shift;
-  Rotate: 2D image rotation;
-  Scale: continuous 2D image downscaling;
-  Euclidean: includes continuous translation, rotation, and reflection,
   forming the group :math:`\mathbb{E}(2)`;
-  Similarity: as above but includes scale, forming the group
   :math:`\text{S}(2)`;
-  Affine: as above but includes shear effects, forming the group
   :math:`\text{Aff}(3)`;
-  Homography: as above but includes perspective (i.e pan and tilt)
   effects, forming the group :math:`\text{PGL}(3)`;
-  PanTiltRotate: pure 3D camera rotation i.e pan, tilt and 2D image
   rotation.

See :ref:`docs <transform>` for full list.

These were proposed in the papers:

-  ``Shift``, ``Rotate``: `Chen et al., Equivariant Imaging: Learning
   Beyond the Range
   Space <https://openaccess.thecvf.com/content/ICCV2021/papers/Chen_Equivariant_Imaging_Learning_Beyond_the_Range_Space_ICCV_2021_paper.pdf>`__
-  ``Scale``: `Scanvic et al., Self-Supervised Learning for Image
   Super-Resolution and Deblurring <https://arxiv.org/abs/2312.11232>`__
-  ``Homography`` and the projective geometry framework: `Wang et al.,
   Perspective-Equivariant Imaging: an Unsupervised Framework for
   Multispectral Pansharpening <https://arxiv.org/abs/2403.09327>`__

.. GENERATED FROM PYTHON SOURCE LINES 37-52

.. code-block:: Python


    import torch
    from torch.utils.data import DataLoader, random_split
    from torchvision.datasets import ImageFolder
    from torchvision.transforms import Compose, ToTensor, CenterCrop, Resize
    from torchvision.datasets.utils import download_and_extract_archive

    import deepinv as dinv
    from deepinv.utils.demo import get_data_home

    device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"

    ORIGINAL_DATA_DIR = get_data_home() / "Urban100"









.. GENERATED FROM PYTHON SOURCE LINES 53-56

Define transforms. For the transforms that involve 3D camera rotation
(i.e pan or tilt), we limit ``theta_max`` for display.


.. GENERATED FROM PYTHON SOURCE LINES 56-69

.. code-block:: Python


    transforms = [
        dinv.transform.Shift(),
        dinv.transform.Rotate(),
        dinv.transform.Scale(),
        dinv.transform.Homography(theta_max=10),
        dinv.transform.projective.Euclidean(),
        dinv.transform.projective.Similarity(),
        dinv.transform.projective.Affine(),
        dinv.transform.projective.PanTiltRotate(theta_max=10),
    ]









.. GENERATED FROM PYTHON SOURCE LINES 70-74

Plot transforms on a sample image. Note that, during training, we never
have access to these ground truth images ``x``, only partial and noisy
measurements ``y``.


.. GENERATED FROM PYTHON SOURCE LINES 74-82

.. code-block:: Python


    x = dinv.utils.load_url_image(dinv.utils.demo.get_image_url("celeba_example.jpg"))
    dinv.utils.plot(
        [x] + [t(x) for t in transforms],
        ["Orig"] + [t.__class__.__name__ for t in transforms],
    )





.. image-sg:: /auto_examples/self-supervised-learning/images/sphx_glr_demo_ei_transforms_001.png
   :alt: Orig, Shift, Rotate, Scale, Homography, Euclidean, Similarity, Affine, PanTiltRotate
   :srcset: /auto_examples/self-supervised-learning/images/sphx_glr_demo_ei_transforms_001.png
   :class: sphx-glr-single-img





.. GENERATED FROM PYTHON SOURCE LINES 83-90

Now, we run an inpainting experiment to reconstruct images from images
masked with a random mask, without ground truth, using EI. For this
example we use the Urban100 images of natural urban scenes. As these
scenes are imaged with a camera free to move and rotate in the world,
all of the above transformations are valid invariances that we can
impose on the unknown image set :math:`x\in X`.


.. GENERATED FROM PYTHON SOURCE LINES 90-106

.. code-block:: Python


    dataset = dinv.datasets.Urban100HR(
        root=ORIGINAL_DATA_DIR,
        download=True,
        transform=Compose([ToTensor(), Resize(256), CenterCrop(256)]),
    )

    train_dataset, test_dataset = random_split(dataset, (0.8, 0.2))

    train_dataloader = DataLoader(train_dataset, shuffle=True)
    test_dataloader = DataLoader(test_dataset)

    # Use physics to generate data online
    physics = dinv.physics.Inpainting((3, 256, 256), mask=0.6, device=device)






.. rst-class:: sphx-glr-script-out

 .. code-block:: none


      0%|          | 0/135388067 [00:00<?, ?it/s]
     18%|█▊        | 23.7M/129M [00:00<00:00, 248MB/s]
     37%|███▋      | 47.4M/129M [00:00<00:00, 244MB/s]
     55%|█████▍    | 70.6M/129M [00:00<00:00, 235MB/s]
     72%|███████▏  | 93.1M/129M [00:00<00:00, 226MB/s]
     89%|████████▉ | 115M/129M [00:00<00:00, 224MB/s] 
    100%|██████████| 129M/129M [00:00<00:00, 231MB/s]

    Extracting:   0%|          | 0/101 [00:00<?, ?it/s]
    Extracting:  19%|█▉        | 19/101 [00:00<00:00, 187.25it/s]
    Extracting:  40%|███▉      | 40/101 [00:00<00:00, 198.89it/s]
    Extracting:  62%|██████▏   | 63/101 [00:00<00:00, 211.14it/s]
    Extracting:  84%|████████▍ | 85/101 [00:00<00:00, 191.86it/s]
    Extracting: 100%|██████████| 101/101 [00:00<00:00, 196.10it/s]
    Dataset has been successfully downloaded.




.. GENERATED FROM PYTHON SOURCE LINES 107-114

For training, use a small UNet, Adam optimizer, EI loss with homography
transform, and the ``deepinv.Trainer`` functionality:

.. note::

      We only train for a single epoch in the demo, but it is recommended to train multiple epochs in practice.


.. GENERATED FROM PYTHON SOURCE LINES 114-142

.. code-block:: Python


    model = dinv.models.UNet(
        in_channels=3, out_channels=3, scales=2, circular_padding=True, batch_norm=False
    ).to(device)

    losses = [
        dinv.loss.MCLoss(),
        dinv.loss.EILoss(dinv.transform.Homography(theta_max=10, device=device)),
    ]

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-8)

    model = dinv.Trainer(
        model=model,
        physics=physics,
        online_measurements=True,
        train_dataloader=train_dataloader,
        eval_dataloader=test_dataloader,
        epochs=1,
        losses=losses,
        optimizer=optimizer,
        verbose=True,
        show_progress_bar=False,
        save_path=None,
        device=device,
    ).train()






.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    The model has 444867 trainable parameters
    Train epoch 0: MCLoss=0.008, EILoss=0.024, TotalLoss=0.032, PSNR=9.754
    Eval epoch 0: PSNR=17.21




.. GENERATED FROM PYTHON SOURCE LINES 143-146

Show results of a pretrained model trained using a larger UNet for 40
epochs:


.. GENERATED FROM PYTHON SOURCE LINES 146-164

.. code-block:: Python


    model = dinv.models.UNet(
        in_channels=3, out_channels=3, scales=3, circular_padding=True, batch_norm=False
    ).to(device)

    ckpt = torch.hub.load_state_dict_from_url(
        dinv.models.utils.get_weights_url("ei", "Urban100_inpainting_homography_model.pth"),
        map_location=device,
    )

    model.load_state_dict(ckpt["state_dict"])

    x = next(iter(train_dataloader))
    x = x.to(device)
    y = physics(x)
    x_hat = model(y)

    dinv.utils.plot([x, y, x_hat], ["x", "y", "reconstruction"])



.. image-sg:: /auto_examples/self-supervised-learning/images/sphx_glr_demo_ei_transforms_002.png
   :alt: x, y, reconstruction
   :srcset: /auto_examples/self-supervised-learning/images/sphx_glr_demo_ei_transforms_002.png
   :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 .. code-block:: none

    Downloading: "https://huggingface.co/deepinv/ei/resolve/main/Urban100_inpainting_homography_model.pth?download=true" to /home/runner/.cache/torch/hub/checkpoints/Urban100_inpainting_homography_model.pth

      0%|          | 0.00/7.90M [00:00<?, ?B/s]
     73%|███████▎  | 5.75M/7.90M [00:00<00:00, 55.2MB/s]
    100%|██████████| 7.90M/7.90M [00:00<00:00, 61.6MB/s]





.. rst-class:: sphx-glr-timing

   **Total running time of the script:** (2 minutes 14.246 seconds)


.. _sphx_glr_download_auto_examples_self-supervised-learning_demo_ei_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_ei_transforms.ipynb <demo_ei_transforms.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: demo_ei_transforms.py <demo_ei_transforms.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: demo_ei_transforms.zip <demo_ei_transforms.zip>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_