Added a whitening transform.

1 files changed, 8 insertions, 0 deletions

diff --git a/src/text_recognizer/datasets/transforms.py b/src/text_recognizer/datasets/transforms.py
index 1105f23..d1ca127 100644
--- a/src/text_recognizer/datasets/transforms.py
+++ b/src/text_recognizer/datasets/transforms.py
@@ -64,3 +64,11 @@ class AddTokens:
             target = torch.cat([sos, target], dim=0)
 
         return target
+
+
+class Whitening:
+    """Whitening of Tensor, i.e. set mean to zero and std to one."""
+
+    def __call__(self, x: Tensor) -> Tensor:
+        """Apply the whitening."""
+        return (x - x.mean()) / x.std()