From 527bb98b191d82b308de1585047e06056258d08d Mon Sep 17 00:00:00 2001 From: aktersnurra 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