Allow v1/v2 independent import of torchvision.transforms

lightly/transforms/__init__.py

@@ -46,6 +46,7 @@
     TiCoView1Transform,
     TiCoView2Transform,
 )
+from lightly.transforms.torchvision_transforms import torchvision_transforms
 from lightly.transforms.vicreg_transform import VICRegTransform, VICRegViewTransform
 from lightly.transforms.vicregl_transform import VICRegLTransform, VICRegLViewTransform
 from lightly.transforms.wmse_transform import WMSETransform

lightly/transforms/torchvision_transforms.py

@@ -0,0 +1,22 @@
+#
+# Copyright (c) Lightly AG and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+#
+from typing import Union
+
+import torch
+from PIL.Image import Image
+from torch import Tensor
+
+try:
+    from torchvision.transforms import v2 as torchvision_transforms
+
+    _TRANSFORMS_V2 = True
+
+except ImportError:
+    from torchvision import transforms as torchvision_transforms
+
+    _TRANSFORMS_V2 = False