use torchvision.transforms.v2

oml/datasets/base.py

@@ -130,7 +130,7 @@ def __init__(
             lru_cache(maxsize=cache_size)(self._read_bytes_image) if cache_size else self._read_bytes_image
         )
 
-        available_augs_types = (albu.Compose, torchvision.transforms.Compose)
+        available_augs_types = (albu.Compose, torchvision.transforms.v2.Compose)
         assert isinstance(self.transform, available_augs_types), f"Type of transforms must be in {available_augs_types}"
 
     @staticmethod
@@ -157,7 +157,7 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
             img = img[y1:y2, x1:x2, :]  # todo: since albu may handle bboxes we should move it to augs
             image_tensor = self.transform(image=img)["image"]
         else:
-            # torchvision.transforms
+            # torchvision.transforms.v2
             img = img.crop((x1, y1, x2, y2))
             image_tensor = self.transform(img)
 

oml/models/vit_unicom/external/vision_transformer.py

@@ -7,7 +7,7 @@
 import torch
 import torch.nn as nn
 from torch.utils.checkpoint import checkpoint
-from torchvision.transforms import (
+from torchvision.transforms.v2 import (
     CenterCrop,
     Compose,
     InterpolationMode,

oml/transforms/images/torchvision.py

@@ -1,5 +1,5 @@
-import torchvision.transforms as t
-from torchvision.transforms import Compose, InterpolationMode, Normalize, ToTensor
+import torchvision.transforms.v2 as t
+from torchvision.transforms.v2 import Compose, InterpolationMode, Normalize, ToTensor
 
 from oml.const import MEAN, STD, TNormParam
 

oml/transforms/images/utils.py

@@ -1,7 +1,7 @@
 from typing import Optional, Union
 
 import albumentations as albu
-import torchvision.transforms as t
+import torchvision.transforms.v2 as t
 
 from oml.utils.images.images import TImReader, imread_cv2, imread_pillow
 

pipelines/README.md

@@ -128,7 +128,7 @@ Let's consider an example of using custom augmentations & extractor to train you
 Your `train.py` and `config.yaml` may look like this:
 ```python
 import hydra
-import torchvision.transforms as t
+import torchvision.transforms.v2 as t
 from omegaconf import DictConfig
 from torchvision.models import resnet18
 

tests/test_runs/test_pipelines/train_with_bboxes.py

@@ -1,5 +1,5 @@
 import hydra
-import torchvision.transforms as t
+import torchvision.transforms.v2 as t
 from omegaconf import DictConfig
 
 from oml.const import HYDRA_BEHAVIOUR, MOCK_DATASET_PATH