oodeel 0.1.1__py3-none-any.whl → 0.3.0__py3-none-any.whl

oodeel/datasets/{ooddataset.py → deprecated/DEPRECATED_ooddataset.py}

@@ -22,11 +22,11 @@
 # SOFTWARE.
 import numpy as np
 
-from ..types import Callable
-from ..types import DatasetType
-from ..types import Optional
-from ..types import Tuple
-from ..types import Union
+from ...types import Callable
+from ...types import DatasetType
+from ...types import Optional
+from ...types import Tuple
+from ...types import Union
 
 
 class OODDataset(object):
@@ -61,6 +61,7 @@ class OODDataset(object):
         load_from_tensorflow_datasets: bool = False,
         input_key: Optional[str] = None,
     ):
+
         self.backend = backend
         self.load_from_tensorflow_datasets = load_from_tensorflow_datasets
 
@@ -74,19 +75,19 @@ class OODDataset(object):
         # Set the channel order depending on the backend
         if self.backend == "torch":
             if load_from_tensorflow_datasets:
-                from .tf_data_handler import TFDataHandler
+                from .DEPRECATED_tf_data_handler import TFDataHandler
                 import tensorflow as tf
 
                 tf.config.set_visible_devices([], "GPU")
                 self._data_handler = TFDataHandler()
                 load_kwargs["as_supervised"] = False
             else:
-                from .torch_data_handler import TorchDataHandler
+                from .DEPRECATED_torch_data_handler import TorchDataHandler
 
                 self._data_handler = TorchDataHandler()
             self.channel_order = "channels_first"
         else:
-            from .tf_data_handler import TFDataHandler
+            from .DEPRECATED_tf_data_handler import TFDataHandler
 
             self._data_handler = TFDataHandler()
             self.channel_order = "channels_last"
@@ -265,7 +266,7 @@ class OODDataset(object):
         with_ood_labels: bool = False,
         with_labels: bool = True,
         shuffle: bool = False,
-        shuffle_buffer_size: Optional[int] = None,
+        **kwargs_prepare,
     ) -> DatasetType:
         """Prepare self.data for scoring or training
 
@@ -282,9 +283,9 @@ class OODDataset(object):
                 Defaults to True.
             shuffle (bool, optional): To shuffle the returned dataset or not.
                 Defaults to False.
-            shuffle_buffer_size (int, optional): (TF only) Size of the shuffle buffer.
-                If None, taken as the number of samples in the dataset.
-                Defaults to None.
+            kwargs_prepare (dict): Additional parameters to be passed to the
+                data_handler.prepare_for_training method.
+
 
         Returns:
             DatasetType: prepared dataset
@@ -323,7 +324,7 @@ class OODDataset(object):
             preprocess_fn=preprocess_fn,
             augment_fn=augment_fn,
             output_keys=keys,
-            shuffle_buffer_size=shuffle_buffer_size,
+            **kwargs_prepare,
         )
 
         return dataset