From a1923e4f33cd3512a3bb3414b257e35865825689 Mon Sep 17 00:00:00 2001
From: Lionel Peer <lionel@lightly.ai>
Date: Thu, 7 Nov 2024 10:43:40 +0100
Subject: [PATCH] Remaining Transforms: Import version independent torchvision
 instead of defaulting to v1 (#1723)

* image grid transform: make v1 independent

* jigsaw: make v1 independent

* MAE: make v1 independent

* MSN: make v1 independent

* multicrop: make v1 independent

* multiview: make v1 independent

* PIRL: make v1 independent

* random crop/flip: make v1 independent

* random rotation: make v1 independent

* simclr transform: make v1 independent

* SimSiam transform: make v1 independent

* SMoG: make v1 independent

* SwaV: make v1 independent

* VICReg: make v1 independent

* VICRegL: make v1 independent

* WMSE: make v1 independent

* fix formatting
---
 lightly/transforms/image_grid_transform.py           | 3 ++-
 lightly/transforms/jigsaw.py                         | 3 ++-
 lightly/transforms/mae_transform.py                  | 2 +-
 lightly/transforms/msn_transform.py                  | 2 +-
 lightly/transforms/multi_crop_transform.py           | 3 +--
 lightly/transforms/multi_view_transform.py           | 3 ++-
 lightly/transforms/pirl_transform.py                 | 3 +--
 lightly/transforms/random_crop_and_flip_with_grid.py | 3 ++-
 lightly/transforms/rotation.py                       | 3 ++-
 lightly/transforms/simclr_transform.py               | 2 +-
 lightly/transforms/simsiam_transform.py              | 2 +-
 lightly/transforms/smog_transform.py                 | 2 +-
 lightly/transforms/swav_transform.py                 | 2 +-
 lightly/transforms/vicreg_transform.py               | 2 +-
 lightly/transforms/vicregl_transform.py              | 2 +-
 lightly/transforms/wmse_transform.py                 | 3 +--
 16 files changed, 21 insertions(+), 19 deletions(-)

diff --git a/lightly/transforms/image_grid_transform.py b/lightly/transforms/image_grid_transform.py
index 1822761a6..7c98880a2 100644
--- a/lightly/transforms/image_grid_transform.py
+++ b/lightly/transforms/image_grid_transform.py
@@ -1,9 +1,10 @@
 from typing import List, Sequence, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
+
 
 class ImageGridTransform:
     """Transforms an image into multiple views and grids.
diff --git a/lightly/transforms/jigsaw.py b/lightly/transforms/jigsaw.py
index 4ed24bb1c..0f81146c3 100644
--- a/lightly/transforms/jigsaw.py
+++ b/lightly/transforms/jigsaw.py
@@ -8,7 +8,8 @@
 from PIL import Image as Image
 from PIL.Image import Image as PILImage
 from torch import Tensor
-from torchvision import transforms as T
+
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 
 if TYPE_CHECKING:
     from numpy.typing import NDArray
diff --git a/lightly/transforms/mae_transform.py b/lightly/transforms/mae_transform.py
index 50f9dd9f7..4f9d89a6b 100644
--- a/lightly/transforms/mae_transform.py
+++ b/lightly/transforms/mae_transform.py
@@ -1,9 +1,9 @@
 from typing import Dict, List, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/msn_transform.py b/lightly/transforms/msn_transform.py
index b8a78ac1c..0f4382c0c 100644
--- a/lightly/transforms/msn_transform.py
+++ b/lightly/transforms/msn_transform.py
@@ -1,11 +1,11 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.multi_view_transform import MultiViewTransform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/multi_crop_transform.py b/lightly/transforms/multi_crop_transform.py
index 49c63229b..a67736996 100644
--- a/lightly/transforms/multi_crop_transform.py
+++ b/lightly/transforms/multi_crop_transform.py
@@ -1,8 +1,7 @@
 from typing import Tuple
 
-import torchvision.transforms as T
-
 from lightly.transforms.multi_view_transform import MultiViewTransform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 
 
 class MultiCropTranform(MultiViewTransform):
diff --git a/lightly/transforms/multi_view_transform.py b/lightly/transforms/multi_view_transform.py
index 62c9f2cb5..7eaef873b 100644
--- a/lightly/transforms/multi_view_transform.py
+++ b/lightly/transforms/multi_view_transform.py
@@ -2,7 +2,8 @@
 
 from PIL.Image import Image
 from torch import Tensor
-from torchvision import transforms as T
+
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 
 
 class MultiViewTransform:
diff --git a/lightly/transforms/pirl_transform.py b/lightly/transforms/pirl_transform.py
index 1d5ec7d57..e62bc3d22 100644
--- a/lightly/transforms/pirl_transform.py
+++ b/lightly/transforms/pirl_transform.py
@@ -1,10 +1,9 @@
 from typing import Dict, List, Tuple, Union
 
-import torchvision.transforms as T
-
 from lightly.transforms.jigsaw import Jigsaw
 from lightly.transforms.multi_view_transform import MultiViewTransform
 from lightly.transforms.rotation import random_rotation_transform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/random_crop_and_flip_with_grid.py b/lightly/transforms/random_crop_and_flip_with_grid.py
index f647c9fee..3df41d4d4 100644
--- a/lightly/transforms/random_crop_and_flip_with_grid.py
+++ b/lightly/transforms/random_crop_and_flip_with_grid.py
@@ -2,11 +2,12 @@
 from typing import Tuple
 
 import torch
-import torchvision.transforms as T
 import torchvision.transforms.functional as F
 from PIL import Image
 from torch import nn
 
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
+
 
 @dataclass
 class Location:
diff --git a/lightly/transforms/rotation.py b/lightly/transforms/rotation.py
index b6b37a63b..fb6ee76d9 100644
--- a/lightly/transforms/rotation.py
+++ b/lightly/transforms/rotation.py
@@ -4,11 +4,12 @@
 from typing import Callable, Tuple, Union
 
 import numpy as np
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 from torchvision.transforms import functional as TF
 
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
+
 
 class RandomRotate:
     """Implementation of random rotation.
diff --git a/lightly/transforms/simclr_transform.py b/lightly/transforms/simclr_transform.py
index 6fd5e7a30..2619f615a 100644
--- a/lightly/transforms/simclr_transform.py
+++ b/lightly/transforms/simclr_transform.py
@@ -1,12 +1,12 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.multi_view_transform import MultiViewTransform
 from lightly.transforms.rotation import random_rotation_transform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/simsiam_transform.py b/lightly/transforms/simsiam_transform.py
index 5121ffc1a..3ed7b1bb1 100644
--- a/lightly/transforms/simsiam_transform.py
+++ b/lightly/transforms/simsiam_transform.py
@@ -1,12 +1,12 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.multi_view_transform import MultiViewTransform
 from lightly.transforms.rotation import random_rotation_transform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/smog_transform.py b/lightly/transforms/smog_transform.py
index 69d8de64c..05697bd99 100644
--- a/lightly/transforms/smog_transform.py
+++ b/lightly/transforms/smog_transform.py
@@ -1,12 +1,12 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.multi_view_transform import MultiViewTransform
 from lightly.transforms.solarize import RandomSolarization
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/swav_transform.py b/lightly/transforms/swav_transform.py
index a34c73647..0cb709a16 100644
--- a/lightly/transforms/swav_transform.py
+++ b/lightly/transforms/swav_transform.py
@@ -1,12 +1,12 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.multi_crop_transform import MultiCropTranform
 from lightly.transforms.rotation import random_rotation_transform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/vicreg_transform.py b/lightly/transforms/vicreg_transform.py
index 1ebe6e04e..5b737b23d 100644
--- a/lightly/transforms/vicreg_transform.py
+++ b/lightly/transforms/vicreg_transform.py
@@ -1,6 +1,5 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
@@ -8,6 +7,7 @@
 from lightly.transforms.multi_view_transform import MultiViewTransform
 from lightly.transforms.rotation import random_rotation_transform
 from lightly.transforms.solarize import RandomSolarization
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/vicregl_transform.py b/lightly/transforms/vicregl_transform.py
index 335c03c34..f48328048 100644
--- a/lightly/transforms/vicregl_transform.py
+++ b/lightly/transforms/vicregl_transform.py
@@ -1,6 +1,5 @@
 from typing import Dict, List, Optional, Tuple, Union
 
-import torchvision.transforms as T
 from PIL.Image import Image
 from torch import Tensor
 
@@ -8,6 +7,7 @@
 from lightly.transforms.image_grid_transform import ImageGridTransform
 from lightly.transforms.random_crop_and_flip_with_grid import RandomResizedCropAndFlip
 from lightly.transforms.solarize import RandomSolarization
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE
 
 
diff --git a/lightly/transforms/wmse_transform.py b/lightly/transforms/wmse_transform.py
index 2d7d2d6dc..8d53d4f44 100644
--- a/lightly/transforms/wmse_transform.py
+++ b/lightly/transforms/wmse_transform.py
@@ -1,9 +1,8 @@
 from typing import Dict, List, Optional, Tuple
 
-import torchvision.transforms as T
-
 from lightly.transforms.gaussian_blur import GaussianBlur
 from lightly.transforms.multi_view_transform import MultiViewTransform
+from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T
 from lightly.transforms.utils import IMAGENET_NORMALIZE