From 38f546f0b86fc0dc89863b00c5ee8c6685394ef2 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Sun, 10 Oct 2021 18:05:24 +0200 Subject: Add custom transforms --- text_recognizer/data/transforms/barlow.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 text_recognizer/data/transforms/barlow.py (limited to 'text_recognizer/data/transforms/barlow.py') diff --git a/text_recognizer/data/transforms/barlow.py b/text_recognizer/data/transforms/barlow.py new file mode 100644 index 0000000..78683cb --- /dev/null +++ b/text_recognizer/data/transforms/barlow.py @@ -0,0 +1,19 @@ +"""Augmentations for training Barlow Twins.""" +from omegaconf.dictconfig import DictConfig +from torch import Tensor + +from text_recognizer.data.transforms.load_transform import load_transform + + +class BarlowTransform: + """Applies two different transforms to input data.""" + + def __init__(self, prim: DictConfig, bis: DictConfig) -> None: + self.prim = load_transform(prim) + self.bis = load_transform(bis) + + def __call__(self, data: Tensor) -> Tensor: + """Applies two different augmentation on the input.""" + x_prim = self.prim(data) + x_bis = self.bis(data) + return x_prim, x_bis -- cgit v1.2.3-70-g09d2