<jupyter_start><jupyter_text>Data Loading<jupyter_code>class dataset(Dataset):

    def __init__(self, image_root, label_root, img_x, img_y):
        """Init function should not do any heavy lifting, but
            must initialize how many items are available in this data set.
        """
        self.images_path = image_root
        self.labels_path = label_root
        self.data_len = 0
        self.images = []
        self.labels = open(self.labels_path, "r").readlines()
        self.transform = transforms.Compose([
            transforms.Resize((img_x, img_y)),  
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
        
        for root, dirs, files in os.walk(self.images_path):
            for file in files:
                if file.endswith('.png'):
                    self.data_len += 1
                    temp = file.split("-")
                    self.images.append(self.images_path + temp[0] + '/' + temp[0] + "-" + temp[1] + "/" + file)

    def __len__(self):
        """return number of points in our dataset"""
        return(self.data_len)

    def __getitem__(self, idx):
        """ Here we have to return the item requested by `idx`
            The PyTorch DataLoader class will use this method to make an iterable for
            our training or validation loop.
        """
        img = self.images[idx]
        label = self.labels[idx]
        img = Image.open(img)
        img = img.convert('RGB')
        img = self.transform(img)
        return(img, label[:-1])
def loader_param():
    img_x = 32
    img_y = 32
    batch_size = 8
    return(img_x, img_y, batch_size)
img_x, img_y, batch_size = loader_param()
train_set = dataset(image_root="./IAM Dataset/words/", label_root = "./IAM Dataset/ascii/labels.txt", img_x = img_x, img_y = img_y)
# train_set.__getitem__(345)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=0)<jupyter_output><empty_output><jupyter_text>Sample visualisation<jupyter_code>train_iter = iter(train_loader)
samples, labels = train_iter.next()
def imshow(img):
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()
trial_num = np.random.randint(0, batch_size)
imshow(samples[trial_num])
labels[trial_num]<jupyter_output><empty_output><jupyter_text>Cuda Support<jupyter_code>device = torch.device("cuda" if torch.cuda.is_available() else "cpu") 
torch.cuda.empty_cache()
# device = "cpu"<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>class Enc_Dec(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        # General Parameters
        self.output_size = output_size
        self.hidden_size = hidden_size
        self.softmax = nn.LogSoftmax(dim = 2)
        
        # Encoder CNN Model
        self.cnn_net = nn.Sequential(
            nn.Conv2d(3, 4, 5, padding = 2), # (N, 3, 32, 128) -> (N, 4, 32, 128)
            nn.ReLU(),
            nn.Conv2d(4, 8, 5, padding = 2), # (N, 4, 32, 128) -> (N, 8, 32, 128)
            nn.ReLU(),
            nn.Conv2d(8, 16, 3, padding = 1), # (N, 8, 32, 128) -> (N, 16, 32, 128)
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, padding = 1), # (N, 16, 32, 128) -> (N, 32, 32, 128)
            nn.ReLU(),
            nn.Conv2d(32, 32, 3, padding = 1) # (N, 32, 32, 128) -> (N, 32, 32, 128)
        )
        
        # Decoder RNN Model
        self.lstm_cell = nn.LSTM(input_size, hidden_size)
        self.h2o = nn.Linear(hidden_size, output_size)
        self.init_hidden = (torch.zeros(1, 1, self.hidden_size, device = device), torch.zeros(1, 1, self.hidden_size, device = device))
    
    def cnn_model(self, input_):       
        return(self.cnn_net(input_))
    
    def rnn_model(self, input_, hidden):
        out, hidden = self.lstm_cell(input_, hidden)
        output = self.h2o(hidden[0])
        output = self.softmax(output)
        return(output, hidden)
    
    def forward(self, input_):
        encoding = self.cnn_model(input_)
#         decodings = []
        decodings_tensor = None
        for i in range(encoding.shape[0]):
#             image_decoding = [] 
            hidden = self.init_hidden
            image_decoding_tensor = None
            for j in range(encoding.shape[3]):
                rnn_input = encoding[i, :, :, j]
                decoding, hidden = self.rnn_model(rnn_input.view(1, 1, -1), hidden)
#                 image_decoding.append(decoding)
                if image_decoding_tensor is None:
                    image_decoding_tensor_size = list(decoding.size())
                    image_decoding_tensor_size.insert(0, encoding.shape[3])
                    image_decoding_tensor = torch.empty(*image_decoding_tensor_size, dtype=decoding.dtype, device=decoding.device)

                image_decoding_tensor[j, :] = decoding   
#             decodings.append(image_decoding)
            if decodings_tensor is None:
                decodings_tensor_size = list(image_decoding_tensor.size())
                decodings_tensor_size.insert(0, encoding.shape[0])
                decodings_tensor = torch.empty(*decodings_tensor_size, dtype=image_decoding_tensor.dtype, device=image_decoding_tensor.device)

            decodings_tensor[i, :] = image_decoding_tensor
        decodings_tensor = torch.squeeze(decodings_tensor)
        decodings_tensor.transpose_(0, 1)
        return(decodings_tensor)
img_txt = Enc_Dec(input_size = 32 * img_x, hidden_size = 256, output_size = 79).to(device)
%%time
decodings = img_txt.forward(input_ = samples.to(device))
decodings.shape<jupyter_output><empty_output><jupyter_text>Train Setup  Data Management<jupyter_code>eng_alphabets = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"!#&\'()*+,-./0123456789:;?'
pad_char = '-PAD-'

eng_alpha2index = {pad_char: 0}
for index, alpha in enumerate(eng_alphabets):
    eng_alpha2index[alpha] = index+1

print(eng_alpha2index)
def word_rep(word, letter2index, max_out_chars, device = 'cpu'):
    rep = torch.zeros(max_out_chars).to(device)
    for letter_index, letter in enumerate(word):
        pos = letter2index[letter]
        rep[letter_index] = pos
    pad_pos = letter2index[pad_char]
    rep[letter_index+1] = pad_pos
    return(rep, len(word))
def words_rep(labels_str, max_out_chars = 20, batch_size = batch_size):
    words_rep = []
    output_cat = None
    lengths_tensor = None
    lengths = []
    for i, label in enumerate(labels_str):
        rep, lnt = word_rep(label, eng_alpha2index, max_out_chars, device)
        words_rep.append(rep)
        if lengths_tensor is None:
            lengths_tensor = torch.empty(len(labels_str), dtype= int)
        if output_cat is None:
            output_cat_size = list(rep.size())
            output_cat_size.insert(0, len(labels_str))
            output_cat = torch.empty(*output_cat_size, dtype=rep.dtype, device=rep.device)
#             print(output_cat.shape)

        output_cat[i, :] = rep
        lengths_tensor[i] = lnt
        lengths.append(lnt)
    return(output_cat, lengths_tensor)<jupyter_output><empty_output><jupyter_text>Training<jupyter_code>def train(epochs = 10, lr = 0.01):
    loss_fn = nn.CTCLoss()
    img_txt = Enc_Dec(input_size = 32 * img_x, hidden_size = 1024, output_size = 79).to(device) 
    optm = optim.SGD(img_txt.parameters(), lr=lr, momentum = 0.9)
    
    input_lengths = torch.full(size=(batch_size,), fill_value=img_y, dtype=torch.long)
    for i in range(epochs):
        total_loss = 0
        for j, (images, labels_str) in enumerate(iter(train_loader)):
            labels, target_lengths = words_rep(labels_str, max_out_chars = 20, batch_size = batch_size)
            images = images.to(device) 
            decodings = img_txt.forward(images)
            print("Label: ", labels[0], '\nOutput: ', torch.argmax(decodings[:, 0, :], 1), '\nTarget len: ', target_lengths[0].item())
#             prob_decodings = nn.functional.log_softmax(decodings, dim = 1)
            loss = loss_fn(decodings, labels, input_lengths, target_lengths)
            
            img_txt.train()
            optm.zero_grad()
            loss.backward()
            optm.step()
            
            torch.cuda.empty_cache()
            total_loss += loss
            print(loss.item(), "\n\n")
        print("Avg Loss for epoch {} is {}\n\n".format(i, total_loss / j))
train()
decodings = img_txt.forward(samples.to(device))
decodings.shape
imshow(samples[0])
torch.argmax(decodings, 1)[:, 0]
rep, _ = words_rep(['Sriram', 'Pingali'])
rep.shape<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12517417556897353, 'fraction_numerical': 0.02333952624245239, 'mean_word_length': 2.6886509635974303, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 30, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '617391', 'license': 'MIT', 'n_tokens_mistral': 3080, 'n_tokens_neox': 2901, 'n_words': 574, 'path': 'Enc_Dec Model.ipynb', 'repo_name': 'SriramPingali/Image-To-Text', 'size': '8612'}