From 4a54d7e690897dd6e6c719fb908fd371a44c2952 Mon Sep 17 00:00:00 2001
From: aktersnurra <gustaf.rydholm@gmail.com>
Date: Sun, 24 Jan 2021 22:14:17 +0100
Subject: Many updates, cool stuff on the way.

---
 src/text_recognizer/datasets/transforms.py | 45 +++++++++++++++++++++++++++++-
 1 file changed, 44 insertions(+), 1 deletion(-)

(limited to 'src/text_recognizer/datasets/transforms.py')

diff --git a/src/text_recognizer/datasets/transforms.py b/src/text_recognizer/datasets/transforms.py
index 8956b01..60987e0 100644
--- a/src/text_recognizer/datasets/transforms.py
+++ b/src/text_recognizer/datasets/transforms.py
@@ -1,14 +1,57 @@
 """Transforms for PyTorch datasets."""
+import random
+
 import numpy as np
 from PIL import Image
 import torch
 from torch import Tensor
 import torch.nn.functional as F
-from torchvision.transforms import Compose, RandomAffine, RandomHorizontalFlip, ToTensor
+from torchvision import transforms
+from torchvision.transforms import (
+    ColorJitter,
+    Compose,
+    Normalize,
+    RandomAffine,
+    RandomHorizontalFlip,
+    RandomRotation,
+    ToPILImage,
+    ToTensor,
+)
 
 from text_recognizer.datasets.util import EmnistMapper
 
 
+class RandomResizeCrop:
+    """Image transform with random resize and crop applied.
+
+    Stolen from
+
+    https://github.com/facebookresearch/gtn_applications/blob/master/datasets/iamdb.py
+
+    """
+
+    def __init__(self, jitter: int = 10, ratio: float = 0.5) -> None:
+        self.jitter = jitter
+        self.ratio = ratio
+
+    def __call__(self, img: np.ndarray) -> np.ndarray:
+        """Applies random crop and rotation to an image."""
+        w, h = img.size
+
+        # pad with white:
+        img = transforms.functional.pad(img, self.jitter, fill=255)
+
+        # crop at random (x, y):
+        x = self.jitter + random.randint(-self.jitter, self.jitter)
+        y = self.jitter + random.randint(-self.jitter, self.jitter)
+
+        # randomize aspect ratio:
+        size_w = w * random.uniform(1 - self.ratio, 1 + self.ratio)
+        size = (h, int(size_w))
+        img = transforms.functional.resized_crop(img, y, x, h, w, size)
+        return img
+
+
 class Transpose:
     """Transposes the EMNIST image to the correct orientation."""
 
-- 
cgit v1.2.3-70-g09d2