Upload 31 files

Model/__init__.py

@@ -0,0 +1 @@
+from Model.model import TTR

Model/attention.py

@@ -0,0 +1,114 @@
+import math
+import torch
+import torch.nn as nn
+
+
+class PositionalEncoding(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)
+    Args:
+        d_model: the embed dim (required).
+        dropout: the dropout value (default=0.1).
+        max_len: the max. length of the incoming sequence (default=5000).
+    Examples:
+        >>> pos_encoder = PositionalEncoding(d_model)
+    """
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        r"""Inputs of forward function
+        Args:
+            x: the sequence fed to the positional encoder model (required).
+        Shape:
+            x: [sequence length, batch size, embed dim]
+            output: [sequence length, batch size, embed dim]
+        Examples:
+            >>> output = pos_encoder(x)
+        """
+
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)
+
+
+def encoder_layer(in_c, out_c, k=3, s=2, p=1):
+    return nn.Sequential(nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c),
+                         nn.ReLU(True))
+
+
+def decoder_layer(in_c, out_c, k=3, s=1, p=1, mode='nearest', scale_factor=None, size=None):
+    align_corners = None if mode == 'nearest' else True
+    return nn.Sequential(nn.Upsample(size=size, scale_factor=scale_factor,
+                                     mode=mode, align_corners=align_corners),
+                         nn.Conv2d(in_c, out_c, k, s, p),
+                         nn.BatchNorm2d(out_c),
+                         nn.ReLU(True))
+
+
+class PositionAttention(nn.Module):
+    def __init__(self, max_length, in_channels=512, num_channels=64,
+                 h=8, w=32, mode='nearest', **kwargs):
+        super().__init__()
+        self.max_length = max_length
+        self.k_encoder = nn.Sequential(
+            encoder_layer(in_channels, num_channels, s=(1, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2)),
+            encoder_layer(num_channels, num_channels, s=(2, 2))
+        )
+        self.k_decoder = nn.Sequential(
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, num_channels, scale_factor=2, mode=mode),
+            decoder_layer(num_channels, in_channels, size=(h, w), mode=mode)
+        )
+
+        self.pos_encoder = PositionalEncoding(in_channels, dropout=0., max_len=max_length)
+        self.project = nn.Linear(in_channels, in_channels)
+
+    def forward(self, x):
+        N, E, H, W = x.size()
+        k, v = x, x  # (N, E, H, W)
+
+        # calculate key vector
+        features = []
+        for i in range(0, len(self.k_encoder)):
+            k = self.k_encoder[i](k)
+            features.append(k)
+        for i in range(0, len(self.k_decoder) - 1):
+            k = self.k_decoder[i](k)
+            k = k + features[len(self.k_decoder) - 2 - i]
+        k = self.k_decoder[-1](k)
+
+        # calculate query vector
+        zeros = x.new_zeros((self.max_length, N, E))  # (T, N, E)
+        q = self.pos_encoder(zeros)  # (T, N, E)
+        q = q.permute(1, 0, 2)  # (N, T, E)
+        q = self.project(q)  # (N, T, E)
+
+        # calculate attention
+        attn_scores = torch.bmm(q, k.flatten(2, 3))  # (N, T, (H*W))
+        attn_scores = attn_scores / (E ** 0.5)
+        attn_scores = torch.softmax(attn_scores, dim=-1)
+
+        v = v.permute(0, 2, 3, 1).view(N, -1, E)  # (N, (H*W), E)
+        attn_vecs = torch.bmm(attn_scores, v)  # (N, T, E)
+
+        return attn_vecs, attn_scores.view(N, -1, H, W)

Model/backbone.py

@@ -0,0 +1,295 @@
+import torch
+import torch.nn as nn
+from timm.models.layers import DropPath, trunc_normal_
+from typing import Optional, Callable
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.ReLU6, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class VITBatchNorm(nn.Module):
+    def __init__(self, num_features):
+        super().__init__()
+        self.num_features = num_features
+        self.bn = nn.BatchNorm1d(num_features=num_features)
+
+    def forward(self, x):
+        return self.bn(x)
+
+
+class Attention(nn.Module):
+    def __init__(self,
+                 dim: int,
+                 num_heads: int = 8,
+                 qkv_bias: bool = False,
+                 qk_scale: Optional[None] = None,
+                 attn_drop: float = 0.,
+                 proj_drop: float = 0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        with torch.cuda.amp.autocast(True):
+            batch_size, num_token, embed_dim = x.shape
+            # qkv is [3,batch_size,num_heads,num_token, embed_dim//num_heads]
+            qkv = self.qkv(x).reshape(
+                batch_size, num_token, 3, self.num_heads, embed_dim // self.num_heads).permute(2, 0, 3, 1, 4)
+        with torch.cuda.amp.autocast(False):
+            q, k, v = qkv[0].float(), qkv[1].float(), qkv[2].float()
+            attn = (q @ k.transpose(-2, -1)) * self.scale
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = (attn @ v).transpose(1, 2).reshape(batch_size, num_token, embed_dim)
+        with torch.cuda.amp.autocast(True):
+            x = self.proj(x)
+            x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 num_heads: int,
+                 num_patches: int,
+                 mlp_ratio: float = 4.,
+                 qkv_bias: bool = False,
+                 qk_scale: Optional[None] = None,
+                 drop: float = 0.,
+                 attn_drop: float = 0.,
+                 drop_path: float = 0.,
+                 act_layer: Callable = nn.ReLU6,
+                 norm_layer: str = "ln",
+                 patch_n: int = 144):
+        super().__init__()
+
+        if norm_layer == "bn":
+            self.norm1 = VITBatchNorm(num_features=num_patches)
+            self.norm2 = VITBatchNorm(num_features=num_patches)
+        elif norm_layer == "ln":
+            self.norm1 = nn.LayerNorm(dim)
+            self.norm2 = nn.LayerNorm(dim)
+
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
+                       act_layer=act_layer, drop=drop)
+        self.extra_gflops = (num_heads * patch_n * (dim // num_heads) * patch_n * 2) / (1000 ** 3)
+
+    def forward(self, x):
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        with torch.cuda.amp.autocast(True):
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    def __init__(self, img_size=[108, 108], patch_size=[9, 9], in_channels=3, embed_dim=768):
+        super().__init__()
+        # img_size = to_2tuple(img_size)
+        # patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * \
+                      (img_size[0] // patch_size[0])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.proj = nn.Conv2d(in_channels, embed_dim,
+                              kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        batch_size, channels, height, width = x.shape
+        assert height == self.img_size[0] and width == self.img_size[1], \
+            f"Input image size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer with support for patch or hybrid CNN input stage
+    """
+
+    def __init__(self,
+                 img_size: list = [112, 112],
+                 patch_size: list = [16, 16],
+                 in_channels: int = 3,
+                 num_classes: int = 1000,
+                 embed_dim: int = 768,
+                 depth: int = 12,
+                 num_heads: int = 12,
+                 mlp_ratio: float = 4.,
+                 qkv_bias: bool = False,
+                 qk_scale: Optional[None] = None,
+                 drop_rate: float = 0.,
+                 attn_drop_rate: float = 0.,
+                 drop_path_rate: float = 0.,
+                 hybrid_backbone: Optional[None] = None,
+                 norm_layer: str = "ln",
+                 mask_ratio=0.1,
+                 using_checkpoint=False,
+                 ):
+        super().__init__()
+        self.num_classes = num_classes
+        # num_features for consistency with other models
+        self.num_features = self.embed_dim = embed_dim
+
+        if hybrid_backbone is not None:
+            raise ValueError
+        else:
+            self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, in_channels=in_channels,
+                                          embed_dim=embed_dim)
+        self.mask_ratio = mask_ratio
+        self.using_checkpoint = using_checkpoint
+        num_patches = self.patch_embed.num_patches
+        self.num_patches = num_patches
+
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+
+        self.patchWSize = img_size[1] // patch_size[1]
+        self.patchHSize = img_size[0] // patch_size[0]
+        patch_n = self.patchWSize * self.patchHSize
+
+        self.blocks = nn.ModuleList(
+            [
+                Block(dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                      drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer,
+                      num_patches=num_patches, patch_n=patch_n)
+                for i in range(depth)]
+        )
+        self.extra_gflops = 0.0
+        for _block in self.blocks:
+            self.extra_gflops += _block.extra_gflops
+
+        if norm_layer == "ln":
+            self.norm = nn.LayerNorm(embed_dim)
+        elif norm_layer == "bn":
+            self.norm = VITBatchNorm(self.num_patches)
+
+        # features head
+        self.feature = nn.Sequential(
+            nn.Linear(in_features=embed_dim * num_patches, out_features=embed_dim, bias=False),
+            nn.BatchNorm1d(num_features=embed_dim, eps=2e-5),
+            nn.Linear(in_features=embed_dim, out_features=num_classes, bias=False),
+            nn.BatchNorm1d(num_features=num_classes, eps=2e-5)
+        )
+
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        torch.nn.init.normal_(self.mask_token, std=.02)
+        trunc_normal_(self.pos_embed, std=.02)
+        # trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {'pos_embed', 'cls_token'}
+
+    def get_classifier(self):
+        return self.head
+
+    def random_masking(self, x, mask_ratio=0.1):
+        """
+        Perform per-sample random masking by per-sample shuffling.
+        Per-sample shuffling is done by argsort random noise.
+        x: [N, L, D], sequence
+        """
+        N, L, D = x.size()  # batch, length, dim
+        len_keep = int(L * (1 - mask_ratio))
+
+        noise = torch.rand(N, L, device=x.device)  # noise in [0, 1]
+
+        # sort noise for each sample
+        # ascend: small is keep, large is remove
+        ids_shuffle = torch.argsort(noise, dim=1)
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset
+        ids_keep = ids_shuffle[:, :len_keep]
+        x_masked = torch.gather(
+            x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, L], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        return x_masked, mask, ids_restore
+
+    def forward_features(self, x):
+        B = x.shape[0]
+        x = self.patch_embed(x)
+        x = x + self.pos_embed
+
+        if self.training and self.mask_ratio > 0:
+            x, _, ids_restore = self.random_masking(x, mask_ratio=self.mask_ratio)
+
+        for func in self.blocks:
+            if self.using_checkpoint and self.training:
+                from torch.utils.checkpoint import checkpoint
+                x = checkpoint(func, x)
+            else:
+                x = func(x)
+        x = self.norm(x.float())
+
+        if self.training and self.mask_ratio > 0:
+            mask_tokens = self.mask_token.repeat(x.shape[0], ids_restore.shape[1] - x.shape[1], 1)
+            x_ = torch.cat([x[:, :, :], mask_tokens], dim=1)  # no cls token
+            x_ = torch.gather(x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2]))  # unshuffle
+            x = x_
+        return x.permute(0, 2, 1).view(B, self.embed_dim, self.patchHSize, self.patchWSize)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+        return x
+
+
+if __name__ == '__main__':
+    img = torch.zeros((1, 3, 32, 128))
+    model = VisionTransformer(img_size=[32, 128],
+                              patch_size=[4, 4],
+                              in_channels=3,
+                              embed_dim=512,
+                              num_heads=8)
+    model.eval()
+    features = model(img)
+    print(features.size())

Model/model.py

@@ -0,0 +1,32 @@
+import torch.nn as nn
+
+from Model.attention import PositionAttention
+from Model.backbone import VisionTransformer
+
+
+class TTR(nn.Module):
+    def __init__(self, args: dict):
+        super().__init__()
+        self.args = args
+
+        self.backbone = VisionTransformer(img_size=args["img_size"],
+                                          patch_size=args["patch_size"],
+                                          in_channels=3,
+                                          embed_dim=args["embed_dim"],
+                                          num_heads=args["num_heads"],
+                                          mask_ratio=args["mask_ratio"])
+
+        self.positionAttention = PositionAttention(max_length=26,
+                                                   in_channels=args["embed_dim"],
+                                                   num_channels=args["position_attention_hidden"],
+                                                   h=args["img_size"][0] // args["patch_size"][0],
+                                                   w=args["img_size"][1] // args["patch_size"][1],
+                                                   mode='nearest')
+        self.cls = nn.Linear(args["embed_dim"], 43)
+        return
+
+    def forward(self, image):
+        features = self.backbone(image)
+        attn_vecs, attn_scores = self.positionAttention(features)
+        logits = self.cls(attn_vecs)
+        return logits

Model/trainer.py

@@ -0,0 +1,264 @@
+import os
+
+import torch
+import tqdm
+from torch import nn
+from torch.nn import functional as F
+
+from torch.utils.data import DataLoader, RandomSampler, Dataset
+
+from metrics import getAcc
+from torch.cuda.amp import autocast, GradScaler
+
+
+class Trainer:
+
+    def __init__(self, args, tb_logger, logger):
+        self.args = args
+
+        self.gpu = torch.device(args.gpu)
+        self.model = None
+        self.it = 0
+        self.best_eval_acc, self.best_it = 0.0, 0
+
+        # init dataset
+        self.trainDataset = None
+        self.trainDataloader = None
+        self.evalDataset = None
+        self.evalDataloader = None
+
+        # optimizer and scheduler
+        self.scheduler = None
+        self.optimizer = None
+
+        # loss
+        self.loss_fn = None
+        self.weight = None
+        self.setLoss(args.loss)
+        self.ignore_index = args.model["letter_size"]
+
+        # gradient clipping
+        if args.clip_grad is not None:
+            self.clip_grad = True
+            self.clip_value = args.clip_grad
+        else:
+            self.clip_grad = False
+
+        if hasattr(args, "label_smoothing") and args.label_smoothing is not None:
+            self.label_smoothing = float(args.label_smoothing)
+        else:
+            self.label_smoothing = 0.0
+
+        # logging
+        if tb_logger is not None:
+            self.tb_log = tb_logger
+        self.print_fn = print if logger is None else logger.info
+
+        return
+
+    def train(self):
+        """
+            Train The Model
+        """
+        self.model.train()
+
+        # for gpu profiling
+        start_batch = torch.cuda.Event(enable_timing=True)
+        end_batch = torch.cuda.Event(enable_timing=True)
+        start_run = torch.cuda.Event(enable_timing=True)
+        end_run = torch.cuda.Event(enable_timing=True)
+
+        scaler = GradScaler()
+
+        start_batch.record()
+        # eval for once
+        if self.args.resume:
+            eval_dict = self.evaluate()
+            print(eval_dict)
+
+        tbar = tqdm.tqdm(total=len(self.trainDataloader), colour='BLUE')
+
+        for samples, targets, _ in self.trainDataloader:
+            tbar.update(1)
+            self.it += 1
+
+            end_batch.record()
+            torch.cuda.synchronize()
+            start_run.record()
+
+            samples, targets = samples.to(self.gpu), targets.to(self.gpu).long()
+
+            with autocast():
+                logits = self.model(samples)
+                loss = F.cross_entropy(logits.flatten(end_dim=1), targets.flatten(),
+                                       ignore_index=self.ignore_index,
+                                       label_smoothing=self.label_smoothing)
+
+            scaler.scale(loss).backward()
+
+            if self.clip_grad:
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip_value)
+
+                scaler.step(self.optimizer)
+                scaler.update()
+                if self.scheduler is not None:
+                    self.scheduler.step()
+                self.model.zero_grad()
+
+            end_run.record()
+            torch.cuda.synchronize()
+
+            # tensorboard_dict update
+            tb_dict = {}
+            tb_dict['train/loss'] = loss.detach().cpu().item()
+
+            tb_dict['lr'] = self.optimizer.param_groups[0]['lr']
+            tb_dict['GPU/prefecth_time'] = start_batch.elapsed_time(end_batch) / 1000.
+            tb_dict['GPU/run_time'] = start_run.elapsed_time(end_run) / 1000.
+
+            if self.it % self.args.num_eval_iter == 0:
+                eval_dict = self.evaluate()
+                tb_dict.update(eval_dict)
+                save_path = self.args.save_path
+                if tb_dict['Word/Acc'] > self.best_eval_acc:
+                    self.best_eval_acc = tb_dict['Word/Acc']
+                    self.best_it = self.it
+
+                self.print_fn(
+                    f"\n {self.it} iteration, {tb_dict}, \n BEST_EVAL_ACC: {self.best_eval_acc}, at {self.best_it} iters")
+                self.print_fn(
+                    f" {self.it} iteration, ACC: {tb_dict['Word/Acc']}\n")
+                if self.it == self.best_it:
+                    self.save_model('model_best.pth', save_path)
+
+            if self.tb_log is not None:
+                self.tb_log.update(tb_dict, self.it)
+            del tb_dict
+            start_batch.record()
+
+        eval_dict = self.evaluate()
+        eval_dict.update({'eval/best_acc': self.best_eval_acc, 'eval/best_it': self.best_it})
+        return eval_dict
+
+    @torch.no_grad()
+    def evaluate(self, model: nn.Module = None, evalDataset: Dataset = None):
+        self.print_fn("\n Evaluation!!!")
+
+        if model is None:
+            model = self.model
+        if evalDataset is not None:
+            evalDataloader = DataLoader(evalDataset, self.args.eval_batch_size, shuffle=False, num_workers=0)
+        else:
+            evalDataloader = self.evalDataloader
+
+        eval_dict = {}
+
+        model.eval()
+
+        preds_arr = None
+        targets_arr = None
+        lengths_arr = None
+        for samples, targets, lengths in evalDataloader:
+            samples, targets = samples.to(self.gpu), targets.to(self.gpu)
+
+            outputs = model(samples)
+
+            preds = torch.max(outputs, dim=2)[1]
+
+            if preds_arr is None:
+                preds_arr = preds.detach().cpu()
+                targets_arr = targets.detach().cpu()
+                lengths_arr = lengths.detach().cpu()
+            else:
+                preds_arr = torch.concat((preds_arr, preds.detach().cpu()))
+                targets_arr = torch.concat((targets_arr, targets.detach().cpu()))
+                lengths_arr = torch.concat((lengths_arr, lengths.detach().cpu()))
+
+        wordAcc, charAcc = getAcc(preds_arr, targets_arr, lengths_arr)
+        eval_dict.update({"Word/Acc": wordAcc,
+                          "Char/Acc": charAcc})
+        model.train()
+        return eval_dict
+
+    def save_model(self, save_name, save_path):
+        save_filename = os.path.join(save_path, save_name)
+        self.model.eval()
+        save_dict = {"model": self.model.state_dict(),
+                     'optimizer': self.optimizer.state_dict(),
+                     'scheduler': self.scheduler.state_dict() if self.scheduler is not None else None,
+                     'it': self.it}
+        torch.save(save_dict, save_filename)
+        self.model.train()
+        self.print_fn(f"model saved: {save_filename}\n")
+
+    def save_baseLearner(self, save_name, save_path, trainIndexes):
+        save_filename = os.path.join(save_path, save_name)
+        self.model.eval()
+        save_dict = {"model": self.model.state_dict(),
+                     'optimizer': self.optimizer.state_dict(),
+                     'scheduler': self.scheduler.state_dict() if self.scheduler is not None else None,
+                     'trainIndexes': trainIndexes,
+                     'it': self.it}
+        torch.save(save_dict, save_filename)
+        self.model.train()
+        self.print_fn(f"model saved: {save_filename}\n")
+
+    def load_model(self, load_dir, load_name):
+        """
+            load saved model a
+        :param load_dir: directory of loading model
+        :param load_name: model name
+        """
+        load_path = os.path.join(load_dir, load_name)
+        checkpoint = torch.load(load_path)
+        self.model.load_state_dict(checkpoint['model'])
+        self.optimizer.load_state_dict(checkpoint['optimizer'])
+        if checkpoint['scheduler'] is not None:
+            self.scheduler.load_state_dict(checkpoint['scheduler'])
+        self.it = checkpoint['it']
+        self.print_fn(f'model loaded from {load_path}')
+
+    def set_optimizer(self, optimizer, scheduler=None):
+        """
+            set optimizer and scheduler
+        :param optimizer: optimizer
+        :param scheduler: scheduler
+        """
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+
+    def setModel(self, model):
+        """
+            set model
+        :param model: model
+        """
+        self.model = model.cuda(self.gpu)
+
+    def setDatasets(self, trainDataset, evalDataset):
+        """
+            set train and evaluation datasets and dataloaders
+        :param trainDataset: train dataset
+        :param evalDataset: evaluation dataset
+        """
+        self.print_fn(f"\n Num Train Labeled Sample : {len(trainDataset)}\n Num Val Sample : {len(evalDataset)}")
+        self.trainDataset = trainDataset
+        self.evalDataset = evalDataset
+
+        self.trainDataloader = DataLoader(trainDataset, batch_size=self.args.batch_size,
+                                          sampler=RandomSampler(data_source=trainDataset,
+                                                                replacement=True,
+                                                                num_samples=self.args.iter * self.args.batch_size),
+                                          num_workers=self.args.num_workers, drop_last=True, pin_memory=True)
+
+        self.evalDataloader = DataLoader(evalDataset, self.args.eval_batch_size, shuffle=False, num_workers=0,
+                                         pin_memory=True)
+
+    def setLoss(self, loss_function: dict):
+        """
+            set loss function
+        :param loss_function: loss function arguments
+        """
+        if loss_function["name"] == 'CrossEntropyLoss':
+            self.loss_fn = nn.CrossEntropyLoss(label_smoothing=loss_function["label_smoothing"]).cuda(self.gpu)
+        else:
+            raise Exception(f"Unknown Loss Function : {loss_function}")

app.py

@@ -0,0 +1,67 @@
+import torch
+from torchvision import transforms
+import gradio as gr
+
+from Model import TTR
+from dataset.charMapper import CharMapper
+
+# arguments
+model_path = "./experiments/real_train/model_best.pth"
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def getTransforms():
+    return transforms.Compose([
+        transforms.Resize((32, 128), transforms.InterpolationMode.BICUBIC),
+        transforms.ToTensor(),
+        transforms.Normalize(0.5, 0.5)
+    ])
+
+
+mapper = CharMapper()
+model = TTR({"img_size": [32, 128],
+             "patch_size": [4, 4],
+             "embed_dim": 512,
+             "num_heads": 8,
+             "position_attention_hidden": 64,
+             "mask_ratio": 0.0
+             })
+model.load_state_dict(torch.load(model_path)["model"])
+model.eval()
+model = model.to("cuda:0")
+
+preprocess = getTransforms()
+
+
+def inference(raw_image):
+    batch = preprocess(raw_image).unsqueeze(0).to(device)
+    outputs = model(batch)
+    preds = torch.max(outputs, dim=2)[1]
+    pred_text = mapper.reverseMapper(preds[0])
+    return pred_text
+
+
+inputs = [gr.Image(type='pil', interactive=True, )]
+outputs = gr.components.Textbox(label="Caption")
+title = "MViT-TR"
+paper_link = "https://www.sciencedirect.com/science/article/pii/S2215098624002672"
+github_link = "https://github.com/serdaryildiz/MViT-TR"
+description = f"<p style='text-align: center'><a href='{github_link}' target='_blank'>MViT-TR</a> : Masked Vision Transformer for Text Recognition"
+examples = [
+    ["fig/0.jpg"],
+    ["fig/145.jpg"],
+    ["fig/195.jpg"],
+    ["fig/270.jpg"],
+]
+article = f"<p style='text-align: center'><a href='{paper_link}' target='_blank'>Paper</a> | <a href='{github_link}' target='_blank'>Github Repo</a></p>"
+css = ".output-image, .input-image, .image-preview {height: 600px !important}"
+
+iface = gr.Interface(fn=inference,
+                     inputs=inputs,
+                     outputs=outputs,
+                     title=title,
+                     description=description,
+                     examples=examples,
+                     article=article,
+                     css=css)
+iface.launch()

dataset/TurkishSceneTextDataset.py

@@ -0,0 +1,70 @@
+import os
+from PIL import Image
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+from dataset.augmentations import CVGeometry, CVDeterioration, CVColorJitter
+from dataset.charMapper import CharMapper
+
+
+class TurkishSceneTextDataset(Dataset):
+    def __init__(self, args: dict, train=True):
+        self.root = args["root"]
+        self.train = train
+
+        txt_path = os.path.join(self.root, "train.txt") if self.train else os.path.join(self.root, "test.txt")
+        with open(txt_path, "r") as fp:
+            lines = fp.readlines()
+
+        self.gt = {}
+        for l in lines:
+            img_name, label = l.strip().split('\t')
+            self.gt[img_name] = label
+
+        self.root = os.path.join(self.root, "train") if self.train else os.path.join(self.root, "test")
+        self.imgPaths = [os.path.join(self.root, p) for p in self.gt.keys()]
+
+        self.transforms = self._getTransforms()
+        self.mapper = CharMapper(letters=args["letters"], maxLength=args["maxLength"])
+
+    def __getitem__(self, item):
+        imgPath = self.imgPaths[item]
+        image = Image.open(imgPath)
+        image = self.transforms(image)
+
+        label = self.gt[os.path.basename(imgPath)]
+
+        label, length = self.mapper(label, return_length=True)
+        return image, label, torch.tensor(length)
+
+    def __len__(self):
+        return len(self.imgPaths)
+
+    def _getTransforms(self):
+        if self.train:
+            return transforms.Compose([
+                CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),
+                CVDeterioration(var=20, degrees=6, factor=4, p=0.25),
+                CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25),
+                transforms.Resize((32, 128), transforms.InterpolationMode.BICUBIC),
+                transforms.ToTensor(),
+                transforms.Normalize(0.5, 0.5)
+            ])
+        else:
+            return transforms.Compose([
+                transforms.Resize((32, 128), transforms.InterpolationMode.BICUBIC),
+                transforms.ToTensor(),
+                transforms.Normalize(0.5, 0.5)
+            ])
+
+
+if __name__ == '__main__':
+    args = {
+        "name": "TurkishSceneTextDataset",
+        "root": "../data/TS-TR",
+        "letters": "0123456789abcçdefgğhıijklmnoöpqrsştuüvwxyz",
+        "maxLength": "25"
+    }
+    dataset = TurkishSceneTextDataset(args, False)

dataset/augmentations.py

@@ -0,0 +1,352 @@
+"""
+    Source : https://github.com/FangShancheng/ABINet/blob/main/transforms.py
+"""
+
+import math
+import numbers
+import random
+
+import cv2
+import numpy as np
+from PIL import Image
+from torchvision import transforms
+from torchvision.transforms import Compose
+
+
+def sample_asym(magnitude, size=None):
+    return np.random.beta(1, 4, size) * magnitude
+
+
+def sample_sym(magnitude, size=None):
+    return (np.random.beta(4, 4, size=size) - 0.5) * 2 * magnitude
+
+
+def sample_uniform(low, high, size=None):
+    return np.random.uniform(low, high, size=size)
+
+
+def get_interpolation(type='random'):
+    if type == 'random':
+        choice = [cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA]
+        interpolation = choice[random.randint(0, len(choice) - 1)]
+    elif type == 'nearest':
+        interpolation = cv2.INTER_NEAREST
+    elif type == 'linear':
+        interpolation = cv2.INTER_LINEAR
+    elif type == 'cubic':
+        interpolation = cv2.INTER_CUBIC
+    elif type == 'area':
+        interpolation = cv2.INTER_AREA
+    else:
+        raise TypeError('Interpolation types only nearest, linear, cubic, area are supported!')
+    return interpolation
+
+
+class CVRandomRotation(object):
+    def __init__(self, degrees=15):
+        assert isinstance(degrees, numbers.Number), "degree should be a single number."
+        assert degrees >= 0, "degree must be positive."
+        self.degrees = degrees
+
+    @staticmethod
+    def get_params(degrees):
+        return sample_sym(degrees)
+
+    def __call__(self, img):
+        angle = self.get_params(self.degrees)
+        src_h, src_w = img.shape[:2]
+        M = cv2.getRotationMatrix2D(center=(src_w / 2, src_h / 2), angle=angle, scale=1.0)
+        abs_cos, abs_sin = abs(M[0, 0]), abs(M[0, 1])
+        dst_w = int(src_h * abs_sin + src_w * abs_cos)
+        dst_h = int(src_h * abs_cos + src_w * abs_sin)
+        M[0, 2] += (dst_w - src_w) / 2
+        M[1, 2] += (dst_h - src_h) / 2
+
+        flags = get_interpolation()
+        return cv2.warpAffine(img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE)
+
+
+class CVRandomAffine(object):
+    def __init__(self, degrees, translate=None, scale=None, shear=None):
+        assert isinstance(degrees, numbers.Number), "degree should be a single number."
+        assert degrees >= 0, "degree must be positive."
+        self.degrees = degrees
+
+        if translate is not None:
+            assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
+                "translate should be a list or tuple and it must be of length 2."
+            for t in translate:
+                if not (0.0 <= t <= 1.0):
+                    raise ValueError("translation values should be between 0 and 1")
+        self.translate = translate
+
+        if scale is not None:
+            assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
+                "scale should be a list or tuple and it must be of length 2."
+            for s in scale:
+                if s <= 0:
+                    raise ValueError("scale values should be positive")
+        self.scale = scale
+
+        if shear is not None:
+            if isinstance(shear, numbers.Number):
+                if shear < 0:
+                    raise ValueError("If shear is a single number, it must be positive.")
+                self.shear = [shear]
+            else:
+                assert isinstance(shear, (tuple, list)) and (len(shear) == 2), \
+                    "shear should be a list or tuple and it must be of length 2."
+                self.shear = shear
+        else:
+            self.shear = shear
+
+    def _get_inverse_affine_matrix(self, center, angle, translate, scale, shear):
+        # https://github.com/pytorch/vision/blob/v0.4.0/torchvision/transforms/functional.py#L717
+        from numpy import sin, cos, tan
+
+        if isinstance(shear, numbers.Number):
+            shear = [shear, 0]
+
+        if not isinstance(shear, (tuple, list)) and len(shear) == 2:
+            raise ValueError(
+                "Shear should be a single value or a tuple/list containing " +
+                "two values. Got {}".format(shear))
+
+        rot = math.radians(angle)
+        sx, sy = [math.radians(s) for s in shear]
+
+        cx, cy = center
+        tx, ty = translate
+
+        # RSS without scaling
+        a = cos(rot - sy) / cos(sy)
+        b = -cos(rot - sy) * tan(sx) / cos(sy) - sin(rot)
+        c = sin(rot - sy) / cos(sy)
+        d = -sin(rot - sy) * tan(sx) / cos(sy) + cos(rot)
+
+        # Inverted rotation matrix with scale and shear
+        # det([[a, b], [c, d]]) == 1, since det(rotation) = 1 and det(shear) = 1
+        M = [d, -b, 0,
+             -c, a, 0]
+        M = [x / scale for x in M]
+
+        # Apply inverse of translation and of center translation: RSS^-1 * C^-1 * T^-1
+        M[2] += M[0] * (-cx - tx) + M[1] * (-cy - ty)
+        M[5] += M[3] * (-cx - tx) + M[4] * (-cy - ty)
+
+        # Apply center translation: C * RSS^-1 * C^-1 * T^-1
+        M[2] += cx
+        M[5] += cy
+        return M
+
+    @staticmethod
+    def get_params(degrees, translate, scale_ranges, shears, height):
+        angle = sample_sym(degrees)
+        if translate is not None:
+            max_dx = translate[0] * height
+            max_dy = translate[1] * height
+            translations = (np.round(sample_sym(max_dx)), np.round(sample_sym(max_dy)))
+        else:
+            translations = (0, 0)
+
+        if scale_ranges is not None:
+            scale = sample_uniform(scale_ranges[0], scale_ranges[1])
+        else:
+            scale = 1.0
+
+        if shears is not None:
+            if len(shears) == 1:
+                shear = [sample_sym(shears[0]), 0.]
+            elif len(shears) == 2:
+                shear = [sample_sym(shears[0]), sample_sym(shears[1])]
+        else:
+            shear = 0.0
+
+        return angle, translations, scale, shear
+
+    def __call__(self, img):
+        src_h, src_w = img.shape[:2]
+        angle, translate, scale, shear = self.get_params(
+            self.degrees, self.translate, self.scale, self.shear, src_h)
+
+        M = self._get_inverse_affine_matrix((src_w / 2, src_h / 2), angle, (0, 0), scale, shear)
+        M = np.array(M).reshape(2, 3)
+
+        startpoints = [(0, 0), (src_w - 1, 0), (src_w - 1, src_h - 1), (0, src_h - 1)]
+        project = lambda x, y, a, b, c: int(a * x + b * y + c)
+        endpoints = [(project(x, y, *M[0]), project(x, y, *M[1])) for x, y in startpoints]
+
+        rect = cv2.minAreaRect(np.array(endpoints))
+        bbox = cv2.boxPoints(rect).astype(dtype=np.int)
+        max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
+        min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
+
+        dst_w = int(max_x - min_x)
+        dst_h = int(max_y - min_y)
+        M[0, 2] += (dst_w - src_w) / 2
+        M[1, 2] += (dst_h - src_h) / 2
+
+        # add translate
+        dst_w += int(abs(translate[0]))
+        dst_h += int(abs(translate[1]))
+        if translate[0] < 0: M[0, 2] += abs(translate[0])
+        if translate[1] < 0: M[1, 2] += abs(translate[1])
+
+        flags = get_interpolation()
+        return cv2.warpAffine(img, M, (dst_w, dst_h), flags=flags, borderMode=cv2.BORDER_REPLICATE)
+
+
+class CVRandomPerspective(object):
+    def __init__(self, distortion=0.5):
+        self.distortion = distortion
+
+    def get_params(self, width, height, distortion):
+        offset_h = sample_asym(distortion * height / 2, size=4).astype(dtype=np.int)
+        offset_w = sample_asym(distortion * width / 2, size=4).astype(dtype=np.int)
+        topleft = (offset_w[0], offset_h[0])
+        topright = (width - 1 - offset_w[1], offset_h[1])
+        botright = (width - 1 - offset_w[2], height - 1 - offset_h[2])
+        botleft = (offset_w[3], height - 1 - offset_h[3])
+
+        startpoints = [(0, 0), (width - 1, 0), (width - 1, height - 1), (0, height - 1)]
+        endpoints = [topleft, topright, botright, botleft]
+        return np.array(startpoints, dtype=np.float32), np.array(endpoints, dtype=np.float32)
+
+    def __call__(self, img):
+        height, width = img.shape[:2]
+        startpoints, endpoints = self.get_params(width, height, self.distortion)
+        M = cv2.getPerspectiveTransform(startpoints, endpoints)
+
+        # TODO: more robust way to crop image
+        rect = cv2.minAreaRect(endpoints)
+        bbox = cv2.boxPoints(rect).astype(dtype=np.int)
+        max_x, max_y = bbox[:, 0].max(), bbox[:, 1].max()
+        min_x, min_y = bbox[:, 0].min(), bbox[:, 1].min()
+        min_x, min_y = max(min_x, 0), max(min_y, 0)
+
+        flags = get_interpolation()
+        img = cv2.warpPerspective(img, M, (max_x, max_y), flags=flags, borderMode=cv2.BORDER_REPLICATE)
+        img = img[min_y:, min_x:]
+        return img
+
+
+class CVRescale(object):
+
+    def __init__(self, factor=4, base_size=(128, 512)):
+        """ Define image scales using gaussian pyramid and rescale image to target scale.
+
+        Args:
+            factor: the decayed factor from base size, factor=4 keeps target scale by default.
+            base_size: base size the build the bottom layer of pyramid
+        """
+        if isinstance(factor, numbers.Number):
+            self.factor = round(sample_uniform(0, factor))
+        elif isinstance(factor, (tuple, list)) and len(factor) == 2:
+            self.factor = round(sample_uniform(factor[0], factor[1]))
+        else:
+            raise Exception('factor must be number or list with length 2')
+        # assert factor is valid
+        self.base_h, self.base_w = base_size[:2]
+
+    def __call__(self, img):
+        if self.factor == 0: return img
+        src_h, src_w = img.shape[:2]
+        cur_w, cur_h = self.base_w, self.base_h
+        scale_img = cv2.resize(img, (cur_w, cur_h), interpolation=get_interpolation())
+        for _ in range(self.factor):
+            scale_img = cv2.pyrDown(scale_img)
+        scale_img = cv2.resize(scale_img, (src_w, src_h), interpolation=get_interpolation())
+        return scale_img
+
+
+class CVGaussianNoise(object):
+    def __init__(self, mean=0, var=20):
+        self.mean = mean
+        if isinstance(var, numbers.Number):
+            self.var = max(int(sample_asym(var)), 1)
+        elif isinstance(var, (tuple, list)) and len(var) == 2:
+            self.var = int(sample_uniform(var[0], var[1]))
+        else:
+            raise Exception('degree must be number or list with length 2')
+
+    def __call__(self, img):
+        noise = np.random.normal(self.mean, self.var ** 0.5, img.shape)
+        img = np.clip(img + noise, 0, 255).astype(np.uint8)
+        return img
+
+
+class CVMotionBlur(object):
+    def __init__(self, degrees=12, angle=90):
+        if isinstance(degrees, numbers.Number):
+            self.degree = max(int(sample_asym(degrees)), 1)
+        elif isinstance(degrees, (tuple, list)) and len(degrees) == 2:
+            self.degree = int(sample_uniform(degrees[0], degrees[1]))
+        else:
+            raise Exception('degree must be number or list with length 2')
+        self.angle = sample_uniform(-angle, angle)
+
+    def __call__(self, img):
+        M = cv2.getRotationMatrix2D((self.degree // 2, self.degree // 2), self.angle, 1)
+        motion_blur_kernel = np.zeros((self.degree, self.degree))
+        motion_blur_kernel[self.degree // 2, :] = 1
+        motion_blur_kernel = cv2.warpAffine(motion_blur_kernel, M, (self.degree, self.degree))
+        motion_blur_kernel = motion_blur_kernel / self.degree
+        img = cv2.filter2D(img, -1, motion_blur_kernel)
+        img = np.clip(img, 0, 255).astype(np.uint8)
+        return img
+
+
+class CVGeometry(object):
+    def __init__(self, degrees=15, translate=(0.3, 0.3), scale=(0.5, 2.),
+                 shear=(45, 15), distortion=0.5, p=0.5):
+        self.p = p
+        type_p = random.random()
+        if type_p < 0.33:
+            self.transforms = CVRandomRotation(degrees=degrees)
+        elif type_p < 0.66:
+            self.transforms = CVRandomAffine(degrees=degrees, translate=translate, scale=scale, shear=shear)
+        else:
+            self.transforms = CVRandomPerspective(distortion=distortion)
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            img = np.array(img)
+            return Image.fromarray(self.transforms(img))
+        else:
+            return img
+
+
+class CVDeterioration(object):
+    def __init__(self, var, degrees, factor, p=0.5):
+        self.p = p
+        transforms = []
+        if var is not None:
+            transforms.append(CVGaussianNoise(var=var))
+        if degrees is not None:
+            transforms.append(CVMotionBlur(degrees=degrees))
+        if factor is not None:
+            transforms.append(CVRescale(factor=factor))
+
+        random.shuffle(transforms)
+        transforms = Compose(transforms)
+        self.transforms = transforms
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            img = np.array(img)
+            return Image.fromarray(self.transforms(img))
+        else:
+            return img
+
+
+class CVColorJitter(object):
+    def __init__(self, brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.5):
+        self.p = p
+        self.transforms = transforms.ColorJitter(brightness=brightness, contrast=contrast,
+                                                 saturation=saturation, hue=hue)
+
+    def __call__(self, img):
+        if random.random() < self.p:
+            return self.transforms(img)
+        else:
+            return img

dataset/charMapper.py

@@ -0,0 +1,55 @@
+import re
+import torch
+
+
+class CharMapper:
+    lower2upper = {
+        ord(u"i"): u"İ",
+        ord(u"ı"): u"I"
+    }
+
+    upper2lower = {
+        ord(u"İ"): u"i",
+        ord(u"I"): u"ı"
+    }
+
+    def __init__(self, letters: str = "0123456789abcçdefgğhıijklmnoöpqrsştuüvwxyz", maxLength: int = 25):
+        self.letters = letters
+        self.maxLength = maxLength
+        self.map = {"[END]": 0}
+        self.reverseMap = {0: "[END]"}
+        i = 1
+        for l in self.letters:
+            self.map[l] = i
+            self.reverseMap[i] = l
+            i += 1
+        self.map["[PAD]"] = i
+        self.reverseMap[i] = "[PAD]"
+        return
+
+    def __call__(self, text: str, return_length=False):
+        text = self.text2label(text)
+        length = len(text) + 1
+        mappedText = torch.tensor([self.map[l] for l in text] + [self.map["[END]"]])
+        text = torch.ones((self.maxLength + 1,)) * self.map["[PAD]"]
+        text[:len(mappedText)] = mappedText
+        if return_length:
+            return text, length
+        else:
+            return text
+
+    def reverseMapper(self, label: torch.tensor):
+        label = label.cpu()
+        text = "".join([self.reverseMap[l] for l in label.numpy()])
+        return text.split("[END]")[0]
+
+    def text2label(self, text):
+        text = re.sub('[^0-9a-zA-ZğüşöçıİĞÜŞÖÇ]+', '', text)
+        text = text.translate(self.upper2lower).lower()
+        return text
+
+
+if __name__ == '__main__':
+    mapper = CharMapper()
+    mapped = mapper("!MA-PİŞ$Z")
+    print(mapped)

dataset/strit.py

@@ -0,0 +1,37 @@
+import os
+
+import torch
+from PIL import Image
+
+from torch.utils.data import Dataset
+from torchvision import transforms
+from dataset.charMapper import CharMapper
+
+
+class STRIT(Dataset):
+    def __init__(self, args: dict):
+        self.root = args["root"]
+        self.imgPaths = [os.path.join(self.root, p) for p in os.listdir(self.root)]
+        self.transforms = self._getTransforms()
+
+        self.mapper = CharMapper(letters=args["letters"], maxLength=args["maxLength"])
+
+    def __getitem__(self, item):
+        imgPath = self.imgPaths[item]
+        image = Image.open(imgPath)
+        image = self.transforms(image)
+        label = imgPath.split('_')[-1].split('.')[0]
+        length = torch.tensor(len(label) + 1)
+        label = self.mapper(label)
+        return image, label, length
+
+    def __len__(self):
+        return len(self.imgPaths)
+
+    @staticmethod
+    def _getTransforms():
+        return transforms.Compose([
+            transforms.Resize((32, 128), transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize(0.5, 0.5)
+        ])

dataset/syntheticTurkishStyleText.py

@@ -0,0 +1,51 @@
+from PIL import Image
+import lmdb
+import six
+
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+
+from dataset.augmentations import CVGeometry, CVDeterioration, CVColorJitter
+from dataset.charMapper import CharMapper
+
+
+class SyntheticTurkishStyleText(Dataset):
+    def __init__(self, args: dict):
+        self.root = args["root"]
+        self.args = args
+        self.transforms = self._getTransforms()
+        self.mapper = CharMapper(letters=args["letters"], maxLength=args["maxLength"])
+
+        self.env = lmdb.open(self.root, readonly=True, lock=False, readahead=False, meminit=False)
+
+        with self.env.begin(write=False) as txn:
+            self.keys = list(txn.cursor().iternext(values=False))
+
+    def __getitem__(self, index):
+        key = self.keys[index]
+        with self.env.begin(write=False) as txn:
+            label = key.decode().split("-*-")[0]
+            imgbuf = txn.get(key)
+            buf = six.BytesIO()
+            buf.write(imgbuf)
+            buf.seek(0)
+            image = Image.open(buf).convert("RGB")
+            image = self.transforms(image)
+        length = torch.tensor(len(label) + 1)
+        label = self.mapper(label)
+        return image, label, length
+
+    def __len__(self):
+        return len(self.keys)
+
+    @staticmethod
+    def _getTransforms():
+        return transforms.Compose([
+            CVGeometry(degrees=45, translate=(0.0, 0.0), scale=(0.5, 2.), shear=(45, 15), distortion=0.5, p=0.5),
+            CVDeterioration(var=20, degrees=6, factor=4, p=0.25),
+            CVColorJitter(brightness=0.5, contrast=0.5, saturation=0.5, hue=0.1, p=0.25),
+            transforms.Resize((32, 128), transforms.InterpolationMode.BICUBIC),
+            transforms.ToTensor(),
+            transforms.Normalize(0.5, 0.5)
+        ])

requirements.txt

@@ -0,0 +1,7 @@
+torch==1.12.1
+torchvision==0.12.1
+opencv-python==4.6.0.66
+transformers==4.27.3
+ftfy==6.1.1
+gradio==3.48.0
+gdown==4.6.0