summaryrefslogtreecommitdiff
path: root/text_recognizer/networks/conv_transformer.py
blob: d36162ae123eac9bf6dea568e800284eedce5448 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
"""Base network module."""
from typing import Type

from torch import Tensor, nn

from text_recognizer.networks.transformer.decoder import Decoder


class ConvTransformer(nn.Module):
    """Base transformer network."""

    def __init__(
        self,
        encoder: Type[nn.Module],
        decoder: Decoder,
    ) -> None:
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def encode(self, img: Tensor) -> Tensor:
        """Encodes images to latent representation."""
        return self.encoder(img)

    def decode(self, tokens: Tensor, img_features: Tensor) -> Tensor:
        """Decodes latent images embedding into characters."""
        return self.decoder(tokens, img_features)

    def forward(self, img: Tensor, tokens: Tensor) -> Tensor:
        """Encodes images into token logtis.

        Args:
            img (Tensor): Input image(s).
            tokens (Tensor): token embeddings.

        Shapes:
            - img: :math: `(B, 1, H, W)`
            - tokens: :math: `(B, Sy)`
            - logits: :math: `(B, Sy, C)`

            where B is the batch size, H is the image height, W is the image
            width, Sy the output length, and C is the number of classes.

        Returns:
            Tensor: Sequence of logits.
        """
        img_features = self.encode(img)
        logits = self.decode(tokens, img_features)
        return logits