From f2cd16f340aa11afadb8fa90c29f85ca1b75a600 Mon Sep 17 00:00:00 2001 From: aktersnurra Date: Mon, 16 Nov 2020 20:26:32 +0100 Subject: Added a whitening transform. --- src/text_recognizer/networks/crnn.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) (limited to 'src/text_recognizer/networks/crnn.py') diff --git a/src/text_recognizer/networks/crnn.py b/src/text_recognizer/networks/crnn.py index 9747429..778e232 100644 --- a/src/text_recognizer/networks/crnn.py +++ b/src/text_recognizer/networks/crnn.py @@ -1,4 +1,4 @@ -"""LSTM with CTC for handwritten text recognition within a line.""" +"""CRNN for handwritten text recognition.""" from typing import Dict, Tuple from einops import rearrange, reduce @@ -89,20 +89,22 @@ class ConvolutionalRecurrentNetwork(nn.Module): x = self.backbone(x) - # Avgerage pooling. + # Average pooling. if self.avg_pool: x = reduce(x, "(b t) c h w -> t b c", "mean", b=b, t=t) else: x = rearrange(x, "(b t) h -> t b h", b=b, t=t) else: # Encode the entire image with a CNN, and use the channels as temporal dimension. - b = x.shape[0] x = self.backbone(x) - x = rearrange(x, "b c h w -> c b (h w)", b=b) + x = rearrange(x, "b c h w -> b w c h") + if self.adaptive_pool is not None: + x = self.adaptive_pool(x) + x = x.squeeze(3) # Sequence predictions. x, _ = self.rnn(x) - # Sequence to classifcation layer. + # Sequence to classification layer. x = self.decoder(x) return x -- cgit v1.2.3-70-g09d2