Repository: deepchecks/deepchecks

Issue:
[FEAT][CV] Optional - run on sample of all batches
A simple way to reduce check / suite run time would be doing the .update() step only on a subset of batches. This should be controlled by a n_samples parameter supplied to the run() method of checks and suites. 


Discussion:

Resolving patch:
diff --git a/deepchecks/core/errors.py b/deepchecks/core/errors.py
--- a/deepchecks/core/errors.py
+++ b/deepchecks/core/errors.py
@@ -18,7 +18,8 @@
     'NumberOfFeaturesLimitError',
     'DatasetValidationError',
     'ModelValidationError',
-    'DeepchecksNotImplementedError'
+    'DeepchecksNotImplementedError',
+    'ValidationError'
 ]
 
 
diff --git a/deepchecks/vision/context.py b/deepchecks/vision/context.py
--- a/deepchecks/vision/context.py
+++ b/deepchecks/vision/context.py
@@ -96,6 +96,7 @@ class Context:
         processing unit for use
     random_state : int
         A seed to set for pseudo-random functions
+    n_samples : int, default: None
     """
 
     def __init__(self,
@@ -106,7 +107,8 @@ def __init__(self,
                  scorers: Mapping[str, Metric] = None,
                  scorers_per_class: Mapping[str, Metric] = None,
                  device: Union[str, torch.device, None] = 'cpu',
-                 random_state: int = 42
+                 random_state: int = 42,
+                 n_samples: int = None
                  ):
         # Validations
         if train is None and test is None and model is None:
@@ -128,11 +130,14 @@ def __init__(self,
                         logger.warning('validate_prediction() was not implemented in %s dataset, '
                                        'some checks will not run', dataset_type)
 
-        # Set seeds to if possible (we set it after the validation to keep the seed state)
-        if train and random_state:
-            train.set_seed(random_state)
-        if test and random_state:
-            test.set_seed(random_state)
+        # The copy does 2 things: Sample n_samples if parameter exists, and shuffle the data.
+        # we shuffle because the data in VisionData is set to be sampled in a fixed order (in the init), so if the user
+        # wants to run without random_state we need to forcefully shuffle (to have different results on different runs
+        # from the same VisionData object), and if there is a random_state the shuffle will always have same result
+        if train:
+            train = train.copy(shuffle=True, n_samples=n_samples, random_state=random_state)
+        if test:
+            test = test.copy(shuffle=True, n_samples=n_samples, random_state=random_state)
 
         self._train = train
         self._test = test
diff --git a/deepchecks/vision/vision_data.py b/deepchecks/vision/vision_data.py
--- a/deepchecks/vision/vision_data.py
+++ b/deepchecks/vision/vision_data.py
@@ -9,16 +9,19 @@
 # ----------------------------------------------------------------------------
 #
 """The vision/dataset module containing the vision Dataset class and its functions."""
+# pylint: disable=protected-access
+import random
 from collections import Counter
 from copy import copy
 from abc import abstractmethod
 from enum import Enum
-from typing import Any, Iterable, List, Optional, Dict, TypeVar, Union
+from typing import Any, List, Optional, Dict, TypeVar, Union, Iterator, Sequence
 
 import logging
 import numpy as np
 import torch
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, BatchSampler, Sampler
+
 
 from deepchecks.core.errors import DeepchecksNotImplementedError, DeepchecksValueError, ValidationError
 from deepchecks.vision.utils.image_functions import ImageInfo
@@ -61,7 +64,8 @@ def __init__(self,
                  label_map: Optional[Dict[int, str]] = None,
                  transform_field: Optional[str] = 'transforms'):
 
-        self._data_loader = self._get_data_loader_copy(data_loader)
+        # Create data loader that uses IndicesSequentialSampler, which always return batches in the same order
+        self._data_loader = self._get_data_loader_sequential(data_loader)
 
         self._num_classes = num_classes
         self._label_map = label_map
@@ -81,7 +85,6 @@ def __init__(self,
         self._n_of_samples_per_class = None
         self._task_type = None
         self._has_label = None
-        self._current_seed = None
 
     @abstractmethod
     def get_classes(self, batch_labels: Union[List[torch.Tensor], torch.Tensor]):
@@ -111,7 +114,7 @@ def validate_prediction(self, batch, model, device):
         )
 
     @abstractmethod
-    def batch_to_images(self, batch) -> Iterable[np.ndarray]:
+    def batch_to_images(self, batch) -> Sequence[np.ndarray]:
         """
         Transform a batch of data to images in the accpeted format.
 
@@ -122,7 +125,7 @@ def batch_to_images(self, batch) -> Iterable[np.ndarray]:
 
         Returns
         -------
-        Iterable[np.ndarray]
+        Sequence[np.ndarray]
             List of images in the accepted format. Each image in the iterable must be a [H, W, C] 3D numpy array.
             See notes for more details.
             :func: `batch_to_images` must be implemented in a subclass.
@@ -232,26 +235,42 @@ def get_augmented_dataset(self, aug) -> VD:
         new_dataset_ref.__setattr__(self._transform_field, new_transform)
         return new_vision_data
 
-    def copy(self) -> VD:
-        """Create new copy of this object, with the data-loader and dataset also copied."""
-        new_vision_data = self.__class__(self._data_loader,
-                                         num_classes=self.num_classes,
-                                         label_map=self._label_map,
-                                         transform_field=self._transform_field)
-        if self._current_seed is not None:
-            new_vision_data.set_seed(self._current_seed)
+    def copy(self, n_samples: int = None, shuffle: bool = False, random_state: int = None) -> VD:
+        """Create new copy of this object, with the data-loader and dataset also copied, and altered by the given \
+        parameters.
+
+        Parameters
+        ----------
+        n_samples : int , default: None
+            take only this number of samples to the copied DataLoader. The samples which will be chosen are affected
+            by random_state (fixed random state will return consistent sampels).
+        shuffle : bool, default: False
+            Whether to shuffle the samples order. The shuffle is affected random_state (fixed random state will return
+            consistent order)
+        random_state : int , default: None
+            random_state used for the psuedo-random actions (sampling and shuffling)
+        """
+        new_vision_data = copy(self)
+        new_vision_data._data_loader = self._get_data_loader_copy(self.data_loader,
+                                                                  shuffle=shuffle,
+                                                                  random_state=random_state,
+                                                                  n_samples=n_samples)
+        # If new data is sampled, then needs to remove cached info
+        if n_samples:
+            new_vision_data._n_of_samples_per_class = None
         return new_vision_data
 
     def to_batch(self, *samples):
         """Use the defined collate_fn to transform a few data items to batch format."""
         return self._data_loader.collate_fn(list(samples))
 
-    def set_seed(self, seed: int):
-        """Set seed for data loader."""
-        generator: torch.Generator = self._data_loader.generator
-        if generator is not None and seed is not None:
-            generator.set_state(torch.Generator().manual_seed(seed).get_state())
-            self._current_seed = seed
+    def batch_of_index(self, *indices):
+        """Return batch samples of the given batch indices."""
+        samples = []
+        for i in indices:
+            index_in_dataset = self.data_loader.batch_sampler.sampler.index_at(i)
+            samples.append(self.data_loader.dataset[index_in_dataset])
+        return self.to_batch(*samples)
 
     def validate_shared_label(self, other: VD):
         """Verify presence of shared labels.
@@ -334,8 +353,52 @@ def __len__(self):
         return len(self._data_loader)
 
     @staticmethod
-    def _get_data_loader_copy(data_loader: DataLoader):
-        props = {
+    def _get_data_loader_copy(data_loader: DataLoader, n_samples: int = None, shuffle: bool = False,
+                              random_state: int = None):
+        """Get a copy of DataLoader which is already using IndicesSequentialSampler, altered by the given parameters.
+
+        Parameters
+        ----------
+        data_loader : DataLoader
+            DataLoader to copy
+        n_samples : int , default: None
+            take only this number of samples to the copied DataLoader. The samples which will be chosen are affected
+            by random_state (fixed random state will return consistent sampels).
+        shuffle : bool, default: False
+            Whether to shuffle the samples order. The shuffle is affected random_state (fixed random state will return
+            consistent order)
+        random_state : int , default: None
+            random_state used for the psuedo-random actions (sampling and shuffling)
+        """
+        # Get sampler and copy it indices if it's already IndicesSequentialSampler
+        batch_sampler = data_loader.batch_sampler
+        if isinstance(batch_sampler.sampler, IndicesSequentialSampler):
+            indices = batch_sampler.sampler.indices
+        else:
+            raise DeepchecksValueError('Expected data loader with sample of type IndicesSequentialSampler')
+        # If got number of samples than take random sample
+        if n_samples:
+            size = min(n_samples, len(indices))
+            if random_state is not None:
+                random.seed(random_state)
+            indices = random.sample(indices, size)
+        # Shuffle indices if need
+        if shuffle:
+            if random_state is not None:
+                random.seed(random_state)
+            indices = random.sample(indices, len(indices))
+        # Create new sampler and batch sampler
+        sampler = IndicesSequentialSampler(indices)
+        new_batch_sampler = BatchSampler(sampler, batch_sampler.batch_size, batch_sampler.drop_last)
+
+        props = VisionData._get_data_loader_props(data_loader)
+        props['batch_sampler'] = new_batch_sampler
+        return data_loader.__class__(**props)
+
+    @staticmethod
+    def _get_data_loader_props(data_loader: DataLoader):
+        """Get properties relevant for the copy of a DataLoader."""
+        return {
             'num_workers': data_loader.num_workers,
             'collate_fn': data_loader.collate_fn,
             'pin_memory': data_loader.pin_memory,
@@ -343,21 +406,50 @@ def _get_data_loader_copy(data_loader: DataLoader):
             'worker_init_fn': data_loader.worker_init_fn,
             'prefetch_factor': data_loader.prefetch_factor,
             'persistent_workers': data_loader.persistent_workers,
-            'generator': torch.Generator()
+            'dataset': copy(data_loader.dataset)
         }
-        # Add batch sampler if exists, else sampler
-        if data_loader.batch_sampler is not None:
-            # Can't deepcopy since generator is not pickle-able,
-            # so copying shallowly and then copies also sampler inside
-            batch_sampler = copy(data_loader.batch_sampler)
-            batch_sampler.sampler = copy(batch_sampler.sampler)
-            # Replace generator instance so the copied dataset will not affect the original
-            batch_sampler.sampler.generator = props['generator']
-            props['batch_sampler'] = batch_sampler
-        else:
-            sampler = copy(data_loader.sampler)
-            # Replace generator instance so the copied dataset will not affect the original
-            sampler.generator = props['generator']
-            props['sampler'] = sampler
-        props['dataset'] = copy(data_loader.dataset)
+
+    @staticmethod
+    def _get_data_loader_sequential(data_loader: DataLoader):
+        """Create new DataLoader with sampler of type IndicesSequentialSampler. This makes the data loader have \
+        consistent batches order."""
+        # First set generator seed to make it reproducible
+        if data_loader.generator:
+            data_loader.generator.set_state(torch.Generator().manual_seed(42).get_state())
+        indices = []
+        batch_sampler = data_loader.batch_sampler
+        # Using the batch sampler to get all indices
+        for batch in batch_sampler:
+            indices += batch
+
+        # Create new sampler and batch sampler
+        sampler = IndicesSequentialSampler(indices)
+        new_batch_sampler = BatchSampler(sampler, batch_sampler.batch_size, batch_sampler.drop_last)
+
+        props = VisionData._get_data_loader_props(data_loader)
+        props['batch_sampler'] = new_batch_sampler
         return data_loader.__class__(**props)
+
+
+class IndicesSequentialSampler(Sampler[int]):
+    """Samples elements sequentially from a given list of indices, without replacement.
+
+    Args:
+        indices (sequence): a sequence of indices
+    """
+
+    indices: List[int]
+
+    def __init__(self, indices: List[int]) -> None:
+        super().__init__(None)
+        self.indices = indices
+
+    def __iter__(self) -> Iterator[int]:
+        return iter(self.indices)
+
+    def __len__(self) -> int:
+        return len(self.indices)
+
+    def index_at(self, location):
+        """Return for a given location, the real index value."""
+        return self.indices[location]


Test patch:
diff --git a/deepchecks/vision/checks/distribution/train_test_label_drift.py b/deepchecks/vision/checks/distribution/train_test_label_drift.py
--- a/deepchecks/vision/checks/distribution/train_test_label_drift.py
+++ b/deepchecks/vision/checks/distribution/train_test_label_drift.py
@@ -57,7 +57,7 @@ class TrainTestLabelDrift(TrainTestCheck):
     ----------
     alternative_label_measurements : List[Dict[str, Any]], default: None
         List of measurements. Replaces the default deepchecks measurements.
-        Each measurement is dictionary with keys 'name' (str), 'method' (Callable) and is_continuous (bool),
+        Each measurement is dictionary with keys 'name' (str), 'method' (Callable) and 'is_continuous' (bool),
         representing attributes of said method.
     max_num_categories : int , default: 10
         Only for non-continues measurements. Max number of allowed categories. If there are more,
diff --git a/tests/vision/base/test_context.py b/tests/vision/base/test_context.py
--- a/tests/vision/base/test_context.py
+++ b/tests/vision/base/test_context.py
@@ -27,7 +27,7 @@
 from deepchecks.core.errors import DeepchecksNotSupportedError
 from deepchecks.core.errors import DatasetValidationError
 from deepchecks.vision.base_checks import Context
-
+from deepchecks.vision import ClassificationData, DetectionData
 
 def test_vision_context_initialization_for_classification_task(mnist_dataset_train, mnist_dataset_test,
                                                                trained_mnist, device):
@@ -42,8 +42,8 @@ def test_vision_context_initialization_for_classification_task(mnist_dataset_tra
 
     # Assert
     assert_that(context, has_properties({
-        'train': same_instance(mnist_dataset_train),
-        'test': same_instance(mnist_dataset_test),
+        'train': instance_of(ClassificationData),
+        'test': instance_of(ClassificationData),
         'model': same_instance(trained_mnist),
         'model_name': equal_to('MNIST'),
         'device': all_of(
@@ -66,8 +66,8 @@ def test_vision_context_initialization_for_object_detection_task(coco_train_visi
 
     # Assert
     assert_that(context, has_properties({
-        'train': same_instance(coco_train_visiondata),
-        'test': same_instance(coco_test_visiondata),
+        'train': instance_of(DetectionData),
+        'test': instance_of(DetectionData),
         'model': same_instance(trained_yolov5_object_detection),
         'model_name': equal_to('COCO'),
         'device': all_of(
diff --git a/tests/vision/base/test_vision_data.py b/tests/vision/base/test_vision_data.py
--- a/tests/vision/base/test_vision_data.py
+++ b/tests/vision/base/test_vision_data.py
@@ -9,6 +9,7 @@
 # ----------------------------------------------------------------------------
 #
 #
+import itertools
 import typing as t
 
 import torch
@@ -21,9 +22,12 @@
     has_entries,
     instance_of,
     all_of,
+    contains_exactly
 )
+import albumentations as A
+import imgaug.augmenters as iaa
 
-from deepchecks.core.errors import ValidationError
+from deepchecks.core.errors import ValidationError, DeepchecksValueError
 from deepchecks.vision.classification_data import ClassificationData
 from deepchecks.vision.vision_data import TaskType
 from deepchecks.vision.datasets.classification.mnist import MNISTData
@@ -32,6 +36,7 @@
 from deepchecks.vision.datasets.detection.coco import COCOData
 from deepchecks.vision.detection_data import DetectionData
 from deepchecks.vision.vision_data import VisionData
+from deepchecks.vision.utils.transformations import AlbumentationsTransformations, ImgaugTransformations
 
 
 class SimpleDetectionData(DetectionData):
@@ -180,3 +185,105 @@ def batch_collate(batch):
     # Act
     # it must not raise an error
     first_dataset.validate_shared_label(second_dataset)
+
+
+def test_get_transforms_type_albumentations(mnist_dataset_train):
+    # Act
+    transform_handler = mnist_dataset_train.get_transform_type()
+    # Assert
+    assert_that(transform_handler, instance_of(AlbumentationsTransformations.__class__))
+
+
+def test_add_augmentation_albumentations(mnist_dataset_train: VisionData):
+    # Arrange
+    augmentation = A.CenterCrop(1, 1)
+    # Act
+    copy_dataset = mnist_dataset_train.get_augmented_dataset(augmentation)
+    # Assert
+    batch = next(iter(copy_dataset.data_loader))
+    data_sample = batch[0][0]
+    assert_that(data_sample.numpy().shape, equal_to((1, 1, 1)))
+
+
+def test_add_augmentation_albumentations_wrong_type(mnist_dataset_train):
+    # Arrange
+    copy_dataset = mnist_dataset_train.copy()
+    augmentation = iaa.CenterCropToFixedSize(1, 1)
+    # Act & Assert
+    msg = r'Transforms is of type albumentations, can\'t add to it type CenterCropToFixedSize'
+    assert_that(calling(copy_dataset.get_augmented_dataset).with_args(augmentation),
+                raises(DeepchecksValueError, msg))
+
+
+def test_get_transforms_type_imgaug(mnist_dataset_train_imgaug):
+    # Act
+    transform_handler = mnist_dataset_train_imgaug.get_transform_type()
+    # Assert
+    assert_that(transform_handler, instance_of(ImgaugTransformations.__class__))
+
+
+def test_add_augmentation_imgaug(mnist_dataset_train_imgaug: VisionData):
+    # Arrange
+    augmentation = iaa.CenterCropToFixedSize(1, 1)
+    # Act
+    copy_dataset = mnist_dataset_train_imgaug.get_augmented_dataset(augmentation)
+    # Assert
+    batch = next(iter(copy_dataset.data_loader))
+    data_sample = batch[0][0]
+    assert_that(data_sample.numpy().shape, equal_to((1, 1, 1)))
+
+
+def test_add_augmentation_imgaug_wrong_type(mnist_dataset_train_imgaug: VisionData):
+    # Arrange
+    copy_dataset = mnist_dataset_train_imgaug.copy()
+    augmentation = A.CenterCrop(1, 1)
+    # Act & Assert
+    msg = r'Transforms is of type imgaug, can\'t add to it type CenterCrop'
+    assert_that(calling(copy_dataset.get_augmented_dataset).with_args(augmentation),
+                raises(DeepchecksValueError, msg))
+
+
+def test_transforms_field_not_exists(mnist_data_loader_train):
+    # Arrange
+    data = VisionData(mnist_data_loader_train, transform_field='not_exists')
+    # Act & Assert
+    msg = r'Underlying Dataset instance does not contain "not_exists" attribute\. If your transformations field is ' \
+          r'named otherwise, you cat set it by using "transform_field" parameter'
+    assert_that(calling(data.get_transform_type).with_args(),
+                raises(DeepchecksValueError, msg))
+
+
+def test_sampler(mnist_dataset_train):
+    # Act
+    sampled = mnist_dataset_train.copy(n_samples=10, random_state=0)
+    # Assert
+    classes = list(itertools.chain(*[b[1].tolist() for b in sampled]))
+    assert_that(classes, contains_exactly(7, 7, 9, 6, 3, 2, 3, 7, 2, 7))
+
+    # Act
+    sampled = mnist_dataset_train.copy(n_samples=500, random_state=0)
+    # Assert
+    total = sum([len(b[0]) for b in sampled])
+    assert_that(total, equal_to(500))
+
+
+def test_data_at_batch_of_index(mnist_dataset_train):
+    # Arrange
+    samples_index = 100
+
+    i = 0
+    for data, labels in mnist_dataset_train.data_loader:
+        if i + len(data) >= samples_index:
+            single_data = data[samples_index - i]
+            single_label = labels[samples_index - i]
+            single_batch = mnist_dataset_train.to_batch((single_data, single_label))
+            break
+        else:
+            i += len(data)
+
+    # Act
+    batch = mnist_dataset_train.batch_of_index(samples_index)
+
+    # Assert
+    assert torch.equal(batch[0], single_batch[0])
+    assert torch.equal(batch[1], single_batch[1])
diff --git a/tests/vision/checks/distribution/heatmap_comparison_test.py b/tests/vision/checks/distribution/heatmap_comparison_test.py
--- a/tests/vision/checks/distribution/heatmap_comparison_test.py
+++ b/tests/vision/checks/distribution/heatmap_comparison_test.py
@@ -55,7 +55,7 @@ def test_object_detection_limit_classes(coco_train_visiondata, coco_test_visiond
     # Assert
     brightness_diff = result.value["diff"]
     assert_that(brightness_diff.mean(), less_than(11))
-    assert_that(brightness_diff.max(), less_than(45))
+    assert_that(brightness_diff.max(), less_than(46))
 
     bbox_diff = result.value["diff_bbox"]
     assert_that(bbox_diff.mean(), less_than(6))
diff --git a/tests/vision/checks/distribution/image_dataset_drift_test.py b/tests/vision/checks/distribution/image_dataset_drift_test.py
--- a/tests/vision/checks/distribution/image_dataset_drift_test.py
+++ b/tests/vision/checks/distribution/image_dataset_drift_test.py
@@ -19,10 +19,9 @@
 
 def add_brightness(img):
     reverse = 255 - img
-    addition_of_brightness = (reverse * 0.1).astype('uint8')
+    addition_of_brightness = (reverse * 0.31).astype(int)
     return img + addition_of_brightness
 
-
 def pil_drift_formatter(batch):
     return [add_brightness(np.array(img)) for img in batch[0]]
 
@@ -36,11 +35,10 @@ def test_no_drift_grayscale(mnist_dataset_train, device):
     result = check.run(train, test, random_state=42, device=device)
     # Assert
     assert_that(result.value, has_entries({
-        'domain_classifier_auc': close_to(0.497, 0.001),
+        'domain_classifier_auc': close_to(0.481, 0.001),
         'domain_classifier_drift_score': equal_to(0),
         'domain_classifier_feature_importance': has_entries({
             'brightness': equal_to(0),
-            'rms_contrast': equal_to(1),
             'aspect_ratio': equal_to(0),
             'area': equal_to(0),
             'normalized_red_mean': equal_to(0),
@@ -59,11 +57,11 @@ def test_drift_grayscale(mnist_dataset_train, mnist_dataset_test, device):
     result = check.run(train, test, random_state=42, device=device)
     # Assert
     assert_that(result.value, has_entries({
-        'domain_classifier_auc': close_to(0.505, 0.001),
-        'domain_classifier_drift_score': close_to(0.011, 0.001),
+        'domain_classifier_auc': close_to(0.524, 0.001),
+        'domain_classifier_drift_score': close_to(0.049, 0.001),
         'domain_classifier_feature_importance': has_entries({
-            'brightness': equal_to(0),
-            'rms_contrast': equal_to(1),
+            'rms_contrast': close_to(0.796, 0.001),
+            'brightness': close_to(0.203, 0.001),
             'aspect_ratio': equal_to(0),
             'area': equal_to(0),
             'normalized_red_mean': equal_to(0),
@@ -83,15 +81,16 @@ def test_no_drift_rgb(coco_train_dataloader, coco_test_dataloader, device):
     result = check.run(train, test, random_state=42, device=device)
     # Assert
     assert_that(result.value, has_entries({
-        'domain_classifier_auc': close_to(0.486, 0.001),
-        'domain_classifier_drift_score': equal_to(0),
+        'domain_classifier_auc': close_to(0.511, 0.001),
+        'domain_classifier_drift_score': close_to(0.023, 0.001),
         'domain_classifier_feature_importance': has_entries({
-            'brightness': equal_to(1),
-            'aspect_ratio': equal_to(0),
-            'area': equal_to(0),
+            'rms_contrast': equal_to(0),
+            'brightness': close_to(0, 0.01),
+            'aspect_ratio': close_to(0, 0.01),
+            'area': close_to(0, 0.001),
             'normalized_red_mean': equal_to(0),
-            'normalized_green_mean': equal_to(0),
-            'normalized_blue_mean': equal_to(0),
+            'normalized_green_mean': close_to(1, 0.001),
+            'normalized_blue_mean': close_to(0, 0.001),
         })
     }))
 
@@ -106,20 +105,18 @@ def batch_to_images(self, batch):
     test = COCOData(coco_test_dataloader)
 
     check = ImageDatasetDrift()
-
     # Act
     result = check.run(train, test, random_state=42, device=device)
     # Assert
     assert_that(result.value, has_entries({
-        'domain_classifier_auc': close_to(0.79, 0.001),
-        'domain_classifier_drift_score': close_to(0.581, 0.001),
+        'domain_classifier_auc': close_to(1, 0.001),
+        'domain_classifier_drift_score': close_to(1, 0.001),
         'domain_classifier_feature_importance': has_entries({
-            'brightness': close_to(0.92, 0.01),
-            'rms_contrast': close_to(0.07, 0.01),
+            'brightness': close_to(1, 0.001),
             'aspect_ratio': equal_to(0),
             'area': equal_to(0),
-            'normalized_red_mean': equal_to(0),
-            'normalized_green_mean': equal_to(0),
-            'normalized_blue_mean': equal_to(0),
+            'normalized_red_mean': close_to(0, 0.01),
+            'normalized_green_mean': close_to(0, 0.01),
+            'normalized_blue_mean': close_to(0, 0.01),
         })
     }))
diff --git a/tests/vision/checks/distribution/simple_feature_contribution_test.py b/tests/vision/checks/distribution/simple_feature_contribution_test.py
--- a/tests/vision/checks/distribution/simple_feature_contribution_test.py
+++ b/tests/vision/checks/distribution/simple_feature_contribution_test.py
@@ -75,10 +75,9 @@ def test_drift_classification(mnist_dataset_train, mnist_dataset_test):
     # Assert
     assert_that(result.value, has_entries({
         'train': has_entries({'brightness': equal_to(0)}),
-        'test': has_entries({'brightness': close_to(0.43, 0.01)}),
-        'train-test difference': has_entries({'brightness': close_to(-0.43, 0.01)})
-    })
-                )
+        'test': has_entries({'brightness': close_to(0.462, 0.001)}),
+        'train-test difference': has_entries({'brightness': close_to(-0.462, 0.001)})
+    }))
 
 
 def test_no_drift_object_detection(coco_train_visiondata):
@@ -94,8 +93,7 @@ def test_no_drift_object_detection(coco_train_visiondata):
         'train': has_entries({'brightness': equal_to(0)}),
         'test': has_entries({'brightness': equal_to(0)}),
         'train-test difference': has_entries({'brightness': equal_to(0)}),
-    })
-                )
+    }))
 
 
 def test_drift_object_detection(coco_train_visiondata, coco_test_visiondata):
@@ -153,6 +151,6 @@ def test_train_test_condition_pps_train_fail(coco_train_visiondata, coco_test_vi
     assert_that(condition_result, equal_condition_result(
         is_pass=False,
         name=f'Train properties\' Predictive Power Score is not greater than {condition_value}',
-        details='Features in train dataset with PPS above threshold: {\'brightness\': \'0.34\', '
+        details='Features in train dataset with PPS above threshold: {\'brightness\': \'0.36\', '
                 '\'rms_contrast\': \'0.34\'}'
     ))
diff --git a/tests/vision/checks/performance/robustness_report_test.py b/tests/vision/checks/performance/robustness_report_test.py
--- a/tests/vision/checks/performance/robustness_report_test.py
+++ b/tests/vision/checks/performance/robustness_report_test.py
@@ -40,11 +40,11 @@ def test_mnist(mnist_dataset_train, trained_mnist, device):
     assert_that(result.value, has_entries({
         'Random Brightness Contrast': has_entries({
             'Precision': has_entries(score=close_to(0.984, 0.001), diff=close_to(-0.002, 0.001)),
-            'Recall': has_entries(score=close_to(0.987, 0.001), diff=close_to(-0.0003, 0.001))
+            'Recall': has_entries(score=close_to(0.986, 0.001), diff=close_to(-0.000, 0.001))
         }),
         'Shift Scale Rotate': has_entries({
-            'Precision': has_entries(score=close_to(0.794, 0.001), diff=close_to(-0.194, 0.001)),
-            'Recall': has_entries(score=close_to(0.776, 0.001), diff=close_to(-0.213, 0.001))
+            'Precision': has_entries(score=close_to(0.800, 0.001), diff=close_to(-0.188, 0.001)),
+            'Recall': has_entries(score=close_to(0.773, 0.001), diff=close_to(-0.216, 0.001))
         }),
     }))
 
@@ -67,8 +67,8 @@ def test_coco_and_condition(coco_train_visiondata, trained_yolov5_object_detecti
     # Assert
     assert_that(result.value, has_entries({
         'Hue Saturation Value': has_entries({
-            'AP': has_entries(score=close_to(0.299, 0.001), diff=close_to(-0.079, 0.001)),
-            'AR': has_entries(score=close_to(0.316, 0.001), diff=close_to(-0.136, 0.001))
+            'AP': has_entries(score=close_to(0.303, 0.001), diff=close_to(-0.064, 0.001)),
+            'AR': has_entries(score=close_to(0.326, 0.001), diff=close_to(-0.109, 0.001))
         }),
     }))
     assert_that(result.conditions_results, has_items(
diff --git a/tests/vision/core/vision_data_test.py b/tests/vision/core/vision_data_test.py
deleted file mode 100644
--- a/tests/vision/core/vision_data_test.py
+++ /dev/null
@@ -1,85 +0,0 @@
-# ----------------------------------------------------------------------------
-# Copyright (C) 2021 Deepchecks (https://www.deepchecks.com)
-#
-# This file is part of Deepchecks.
-# Deepchecks is distributed under the terms of the GNU Affero General
-# Public License (version 3 or later).
-# You should have received a copy of the GNU Affero General Public License
-# along with Deepchecks.  If not, see <http://www.gnu.org/licenses/>.
-# ----------------------------------------------------------------------------
-#
-import imgaug.augmenters as iaa
-import albumentations as A
-
-from deepchecks.core.errors import DeepchecksValueError
-from deepchecks.vision.utils.transformations import AlbumentationsTransformations, ImgaugTransformations
-from hamcrest import assert_that, instance_of, equal_to, calling, raises
-
-from tests.vision.vision_conftest import *
-from deepchecks.vision import VisionData
-
-
-def test_get_transforms_type_albumentations(mnist_dataset_train):
-    # Act
-    transform_handler = mnist_dataset_train.get_transform_type()
-    # Assert
-    assert_that(transform_handler, instance_of(AlbumentationsTransformations.__class__))
-
-
-def test_add_augmentation_albumentations(mnist_dataset_train: VisionData):
-    # Arrange
-    augmentation = A.CenterCrop(1, 1)
-    # Act
-    copy_dataset = mnist_dataset_train.get_augmented_dataset(augmentation)
-    # Assert
-    batch = next(iter(copy_dataset.data_loader))
-    data_sample = batch[0][0]
-    assert_that(data_sample.numpy().shape, equal_to((1, 1, 1)))
-
-
-def test_add_augmentation_albumentations_wrong_type(mnist_dataset_train):
-    # Arrange
-    copy_dataset = mnist_dataset_train.copy()
-    augmentation = iaa.CenterCropToFixedSize(1, 1)
-    # Act & Assert
-    msg = r'Transforms is of type albumentations, can\'t add to it type CenterCropToFixedSize'
-    assert_that(calling(copy_dataset.get_augmented_dataset).with_args(augmentation),
-                raises(DeepchecksValueError, msg))
-
-
-def test_get_transforms_type_imgaug(mnist_dataset_train_imgaug):
-    # Act
-    transform_handler = mnist_dataset_train_imgaug.get_transform_type()
-    # Assert
-    assert_that(transform_handler, instance_of(ImgaugTransformations.__class__))
-
-
-def test_add_augmentation_imgaug(mnist_dataset_train_imgaug: VisionData):
-    # Arrange
-    augmentation = iaa.CenterCropToFixedSize(1, 1)
-    # Act
-    copy_dataset = mnist_dataset_train_imgaug.get_augmented_dataset(augmentation)
-    # Assert
-    batch = next(iter(copy_dataset.data_loader))
-    data_sample = batch[0][0]
-    assert_that(data_sample.numpy().shape, equal_to((1, 1, 1)))
-
-
-def test_add_augmentation_imgaug_wrong_type(mnist_dataset_train_imgaug: VisionData):
-    # Arrange
-    copy_dataset = mnist_dataset_train_imgaug.copy()
-    augmentation = A.CenterCrop(1, 1)
-    # Act & Assert
-    msg = r'Transforms is of type imgaug, can\'t add to it type CenterCrop'
-    assert_that(calling(copy_dataset.get_augmented_dataset).with_args(augmentation),
-                raises(DeepchecksValueError, msg))
-
-
-def test_transforms_field_not_exists(mnist_data_loader_train):
-    # Arrange
-    data = VisionData(mnist_data_loader_train, transform_field='not_exists')
-    # Act & Assert
-    msg = r'Underlying Dataset instance does not contain "not_exists" attribute\. If your transformations field is ' \
-          r'named otherwise, you cat set it by using "transform_field" parameter'
-    assert_that(calling(data.get_transform_type).with_args(),
-                raises(DeepchecksValueError, msg))