From 527bb98b191d82b308de1585047e06056258d08d Mon Sep 17 00:00:00 2001
From: aktersnurra <gustaf.rydholm@gmail.com>
Date: Wed, 18 Nov 2020 20:56:19 +0100
Subject: Some minor changes.

---
 src/text_recognizer/datasets/transforms.py | 15 ++++++++++-----
 1 file changed, 10 insertions(+), 5 deletions(-)

(limited to 'src/text_recognizer/datasets')

diff --git a/src/text_recognizer/datasets/transforms.py b/src/text_recognizer/datasets/transforms.py
index d1ca127..1ec23dc 100644
--- a/src/text_recognizer/datasets/transforms.py
+++ b/src/text_recognizer/datasets/transforms.py
@@ -4,7 +4,7 @@ from PIL import Image
 import torch
 from torch import Tensor
 import torch.nn.functional as F
-from torchvision.transforms import Compose, ToPILImage, ToTensor
+from torchvision.transforms import Compose, RandomAffine, ToTensor
 
 from text_recognizer.datasets.util import EmnistMapper
 
@@ -66,9 +66,14 @@ class AddTokens:
         return target
 
 
-class Whitening:
-    """Whitening of Tensor, i.e. set mean to zero and std to one."""
+class ApplyContrast:
+    """Sets everything below a threshold to zero, i.e. increase contrast."""
+
+    def __init__(self, low: float = 0.0, high: float = 0.25) -> None:
+        self.low = low
+        self.high = high
 
     def __call__(self, x: Tensor) -> Tensor:
-        """Apply the whitening."""
-        return (x - x.mean()) / x.std()
+        """Apply mask binary mask to input tensor."""
+        mask = x > np.random.RandomState().uniform(low=self.low, high=self.high)
+        return x * mask
-- 
cgit v1.2.3-70-g09d2