.. 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

        New to DeepInverse? Get started with the basics with the
        :ref:`5 minute quickstart tutorial <sphx_glr_auto_examples_basics_demo_quickstart.py>`.

.. 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``: :footcite:t:`chen2021equivariant`.
-  ``Scale``: :footcite:t:`scanvic2025scale`.
-  ``Homography`` and the projective geometry framework: :footcite:t:`wang2024perspective`.

.. GENERATED FROM PYTHON SOURCE LINES 32-45

.. code-block:: Python


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

    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 46-49

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 49-62

.. 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),
    ]






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

 .. code-block:: none

    /home/runner/work/deepinv/deepinv/deepinv/transform/rotate.py:46: 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 63-67

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 67-76

.. code-block:: Python


    x = dinv.utils.load_example("celeba_example.jpg")
    dinv.utils.plot(
        [x] + [t(x) for t in transforms],
        ["Orig"] + [t.__class__.__name__ for t in transforms],
        fontsize=24,
    )





.. 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 77-84

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 84-100

.. 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]
      3%|▎         | 3.38M/129M [00:00<00:03, 35.3MB/s]
     23%|██▎       | 30.3M/129M [00:00<00:00, 180MB/s] 
     45%|████▍     | 57.8M/129M [00:00<00:00, 230MB/s]
     66%|██████▋   | 85.8M/129M [00:00<00:00, 255MB/s]
     88%|████████▊ | 114M/129M [00:00<00:00, 268MB/s] 
    100%|██████████| 129M/129M [00:00<00:00, 243MB/s]

    Extracting:   0%|          | 0/101 [00:00<?, ?it/s]
    Extracting:  21%|██        | 21/101 [00:00<00:00, 202.52it/s]
    Extracting:  48%|████▊     | 48/101 [00:00<00:00, 239.73it/s]
    Extracting:  72%|███████▏  | 73/101 [00:00<00:00, 223.67it/s]
    Extracting:  95%|█████████▌| 96/101 [00:00<00:00, 220.03it/s]
    Extracting: 100%|██████████| 101/101 [00:00<00:00, 222.99it/s]
    Dataset has been successfully downloaded.




.. GENERATED FROM PYTHON SOURCE LINES 101-108

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 108-136

.. 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.006, EILoss=0.021, TotalLoss=0.028, PSNR=12.142
    Eval epoch 0: PSNR=19.328
    Best model saved at epoch 1




.. GENERATED FROM PYTHON SOURCE LINES 137-140

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


.. GENERATED FROM PYTHON SOURCE LINES 140-159

.. 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]
     98%|█████████▊| 7.75M/7.90M [00:00<00:00, 80.8MB/s]
    100%|██████████| 7.90M/7.90M [00:00<00:00, 81.8MB/s]




.. GENERATED FROM PYTHON SOURCE LINES 160-163

:References:

.. footbibliography::


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

   **Total running time of the script:** (2 minutes 18.447 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>`_