braindecode 0.8.1__py3-none-any.whl → 1.1.0__py3-none-any.whl

braindecode/__init__.py

@@ -1,7 +1,6 @@
-from .version import __version__
-
 from .classifier import EEGClassifier
 from .regressor import EEGRegressor
+from .version import __version__
 
 __all__ = [
     "__version__",

braindecode/augmentation/__init__.py

@@ -1,30 +1,50 @@
 """
 Utilities for data augmentation.
 """
-from .base import Transform, IdentityTransform, Compose, AugmentedDataLoader
+
+from . import functional
+from .base import AugmentedDataLoader, Compose, IdentityTransform, Transform
 from .transforms import (
-    TimeReverse,
-    SignFlip,
-    FTSurrogate,
-    ChannelsShuffle,
+    BandstopFilter,
     ChannelsDropout,
-    GaussianNoise,
+    ChannelsShuffle,
     ChannelsSymmetry,
-    SmoothTimeMask,
-    BandstopFilter,
     FrequencyShift,
+    FTSurrogate,
+    GaussianNoise,
+    MaskEncoding,
+    Mixup,
+    SegmentationReconstruction,
     SensorsRotation,
-    SensorsZRotation,
-    SensorsYRotation,
     SensorsXRotation,
-    Mixup,
+    SensorsYRotation,
+    SensorsZRotation,
+    SignFlip,
+    SmoothTimeMask,
+    TimeReverse,
 )
 
-from . import functional
-
-__all__ = ["Transform", "IdentityTransform", "Compose", "AugmentedDataLoader",
-           "TimeReverse", "SignFlip", "FTSurrogate", "ChannelsShuffle",
-           "ChannelsDropout", "GaussianNoise", "ChannelsSymmetry",
-           "SmoothTimeMask", "BandstopFilter", "FrequencyShift",
-           "SensorsRotation", "SensorsZRotation", "SensorsYRotation",
-           "SensorsXRotation", "Mixup", "functional"]
+__all__ = [
+    "Transform",
+    "IdentityTransform",
+    "Compose",
+    "AugmentedDataLoader",
+    "TimeReverse",
+    "SignFlip",
+    "FTSurrogate",
+    "ChannelsShuffle",
+    "ChannelsDropout",
+    "GaussianNoise",
+    "ChannelsSymmetry",
+    "SmoothTimeMask",
+    "BandstopFilter",
+    "FrequencyShift",
+    "SensorsRotation",
+    "SensorsZRotation",
+    "SensorsYRotation",
+    "SensorsXRotation",
+    "Mixup",
+    "SegmentationReconstruction",
+    "MaskEncoding",
+    "functional",
+]

braindecode/augmentation/base.py

@@ -5,31 +5,28 @@
 #          Valentin Iovene <val@too.gy>
 # License: BSD (3-clause)
 
-from typing import List, Tuple, Any, Optional, Union, Callable
 from numbers import Real
+from typing import Any, Callable, Optional, Union
 
-from sklearn.utils import check_random_state
 import torch
+from sklearn.utils import check_random_state
 from torch import Tensor, nn
 from torch.utils.data import DataLoader
 from torch.utils.data._utils.collate import default_collate
 
 from .functional import identity
 
-Batch = List[Tuple[torch.Tensor, int, Any]]
+Batch = list[tuple[torch.Tensor, int, Any]]
 Output = Union[
     # just outputting X
     torch.Tensor,
     # outputting (X, y) where y can be a tensor or tuple of tensors
-    Tuple[torch.Tensor, Union[torch.Tensor, Tuple[torch.Tensor, ...]]]
+    tuple[torch.Tensor, Union[torch.Tensor, tuple[torch.Tensor, ...]]],
 ]
 # (X, y) -> (X', y') where y' can be a tensor or a tuple of tensors
 Operation = Callable[
     [torch.Tensor, torch.Tensor],
-    Tuple[
-        torch.Tensor,
-        Union[torch.Tensor, Tuple[torch.Tensor, ...]]
-    ]
+    tuple[torch.Tensor, Union[torch.Tensor, tuple[torch.Tensor, ...]]],
 ]
 
 
@@ -50,18 +47,20 @@ class Transform(torch.nn.Module):
         Used to decide whether or not to transform given the probability
         argument. Defaults to None.
     """
+
     operation: Operation
 
     def __init__(self, probability=1.0, random_state=None):
         super().__init__()
         if self.forward.__func__ is Transform.forward:
-            assert callable(self.operation), "operation should be a " \
-                                             "``callable``. "
+            assert callable(self.operation), "operation should be a ``callable``. "
 
         assert isinstance(probability, Real), (
-            f"probability should be a ``real``. Got {type(probability)}.")
-        assert probability <= 1. and probability >= 0., \
+            f"probability should be a ``real``. Got {type(probability)}."
+        )
+        assert probability <= 1.0 and probability >= 0.0, (
             "probability should be between 0 and 1."
+        )
         self._probability = probability
         self.rng = check_random_state(random_state)
 
@@ -108,8 +107,9 @@ class Transform(torch.nn.Module):
         if num_valid > 0:
             # Uses the mask to define the output
             out_X[mask, ...], tr_y = self.operation(
-                out_X[mask, ...], out_y[mask],
-                **self.get_augmentation_params(out_X[mask, ...], out_y[mask])
+                out_X[mask, ...],
+                out_y[mask],
+                **self.get_augmentation_params(out_X[mask, ...], out_y[mask]),
             )
             # Apply the operation defining the Transform to the whole batch
             if isinstance(tr_y, tuple):
@@ -125,11 +125,10 @@ class Transform(torch.nn.Module):
             return out_X
 
     def _get_mask(self, batch_size, device) -> torch.Tensor:
-        """Samples whether to apply operation or not over the whole batch
-        """
-        return torch.as_tensor(
-            self.probability > self.rng.uniform(size=batch_size)
-        ).to(device)
+        """Samples whether to apply operation or not over the whole batch"""
+        return torch.as_tensor(self.probability > self.rng.uniform(size=batch_size)).to(
+            device
+        )
 
     @property
     def probability(self):
@@ -141,7 +140,9 @@ class IdentityTransform(Transform):
 
     Transform that does not change the input.
     """
-    operation = staticmethod(identity)
+
+    operation = staticmethod(identity)  # type: ignore[assignment]
+    # https://github.com/python/mypy/issues/4574
 
 
 class Compose(Transform):
@@ -167,8 +168,8 @@ class Compose(Transform):
 
 
 def _make_collateable(transform, device=None):
-    """ Wraps a transform to make it collateable.
-        with device control. """
+    """Wraps a transform to make it collateable.
+    with device control."""
 
     def _collate_fn(batch):
         collated_batch = default_collate(batch)
@@ -205,7 +206,7 @@ class AugmentedDataLoader(DataLoader):
                 "to pass transform"
             )
         if transforms is None or (
-                isinstance(transforms, list) and len(transforms) == 0
+            isinstance(transforms, list) and len(transforms) == 0
         ):
             self.collated_tr = _make_collateable(IdentityTransform(), device=device)
         elif isinstance(transforms, (Transform, nn.Module)):
@@ -218,8 +219,4 @@ class AugmentedDataLoader(DataLoader):
                 "or a list of Transform objects."
             )
 
-        super().__init__(
-            dataset,
-            collate_fn=self.collated_tr,
-            **kwargs
-        )
+        super().__init__(dataset, collate_fn=self.collated_tr, **kwargs)