from PIL import Image
import os
import torch
from tqdm import tqdm

if __name__ != "__main__":
    from . import utils,dataset,model


class seq2seq_recognize_unit():
    def __init__(self,alphabets_path,encoder_pth,decoder_pth):
        print("Creating seq2seq recognize unit...")
        self.alphabets = utils.get_alphabets(alphabets_path)
        self.converter = utils.ConvertBetweenStringAndLabel(self.alphabets)
        self.num_classes = len(self.alphabets) + 2

        self.encoder = model.Encoder(channel_size=1, hidden_size=256)
        self.decoder = model.Decoder(hidden_size=256, output_size=self.num_classes, dropout_p=0, max_length=71)
        self.encoder.load_state_dict(torch.load(encoder_pth))
        print('load encoder pth %s' % encoder_pth)
        self.decoder.load_state_dict(torch.load(decoder_pth))
        print('load encoder pth %s' % decoder_pth)

        self.encoder.cuda()
        self.decoder.cuda()
        self.encoder.eval()
        self.decoder.eval()

        self.trans = dataset.resizeNormalize(32)

    def image_to_tensor(self,img):
        img = self.trans(img)
        img = torch.cat([img, torch.zeros([1, img.shape[1], 280 - img.shape[2]])], 2).unsqueeze(0)
        return img

    def recognize(self,img_tensor,max_length = 20):
        encoder_out = self.encoder(img_tensor.cuda())
        decoder_input = torch.zeros(1).long().cuda()
        decoder_hidden = self.decoder.initHidden(1).cuda()

        decoded_words = []
        prob = 1.0
        for di in range(max_length):
            decoder_output, decoder_hidden, decoder_attention = self.decoder(decoder_input, decoder_hidden, encoder_out)
            probs = torch.exp(decoder_output)
            _, topi = decoder_output.data.topk(1)
            ni = topi.squeeze(1)
            decoder_input = ni
            prob *= probs[:, ni]
            if ni == utils.EOS_TOKEN:
                break
            else:
                decoded_words.append(self.converter.decode(ni))
        words=''.join(decoded_words)
        #print('predict_string: {} => predict_probility: {}'.format(words, prob.item()))
        return words


if __name__ == "__main__":
    import utils, dataset, model
    alphabets_path = "../datasets/alphabets.txt"
    encoder_pth = "./pths/encoder_150.pth"
    decoder_pth = "./pths/decoder_150.pth"
    input_img_path = "../datasets/CRNN_dataset/test/img/"
    input_gt_path = "../datasets/CRNN_dataset/test/gt/"
    output_path = "../datasets/output/"
    recognize_unit = seq2seq_recognize_unit(alphabets_path,encoder_pth,decoder_pth)
    names =[ i[:-4] for i in os.listdir(input_img_path)]
    for name in tqdm(names):
        img=Image.open(input_img_path+name+".jpg")
        img_tensor = recognize_unit.image_to_tensor(img)
        words = recognize_unit.recognize(img_tensor)
        with open(input_gt_path+name+".txt","r",encoding="utf8") as f:
            box = f.read()
        gt = box+','+words
        #print(gt)
        with open(output_path+name+".txt","w",encoding="utf8") as f:
            f.write(gt)




from PIL import Image
import os
import torch
from tqdm import tqdm

if __name__ != "__main__":
    from . import utils,dataset,model


class seq2seq_recognize_unit():
    def __init__(self,alphabets_path,encoder_pth,decoder_pth):
        print("Creating seq2seq recognize unit...")
        self.alphabets = utils.get_alphabets(alphabets_path)
        self.converter = utils.ConvertBetweenStringAndLabel(self.alphabets)
        self.num_classes = len(self.alphabets) + 2

        self.encoder = model.Encoder(channel_size=1, hidden_size=256)
        self.decoder = model.Decoder(hidden_size=256, output_size=self.num_classes, dropout_p=0, max_length=71)
        self.encoder.load_state_dict(torch.load(encoder_pth))
        print('load encoder pth %s' % encoder_pth)
        self.decoder.load_state_dict(torch.load(decoder_pth))
        print('load encoder pth %s' % decoder_pth)

        self.encoder.cuda()
        self.decoder.cuda()
        self.encoder.eval()
        self.decoder.eval()

        self.trans = dataset.resizeNormalize(32)

    def image_to_tensor(self,img):
        img = self.trans(img)
        img = torch.cat([img, torch.zeros([1, img.shape[1], 280 - img.shape[2]])], 2).unsqueeze(0)
        return img

    def recognize(self,img_tensor,max_length = 20):
        encoder_out = self.encoder(img_tensor.cuda())
        decoder_input = torch.zeros(1).long().cuda()
        decoder_hidden = self.decoder.initHidden(1).cuda()

        decoded_words = []
        prob = 1.0
        for di in range(max_length):
            decoder_output, decoder_hidden, decoder_attention = self.decoder(decoder_input, decoder_hidden, encoder_out)
            probs = torch.exp(decoder_output)
            _, topi = decoder_output.data.topk(1)
            ni = topi.squeeze(1)
            decoder_input = ni
            prob *= probs[:, ni]
            if ni == utils.EOS_TOKEN:
                break
            else:
                decoded_words.append(self.converter.decode(ni))
        words=''.join(decoded_words)
        #print('predict_string: {} => predict_probility: {}'.format(words, prob.item()))
        return words


if __name__ == "__main__":
    import utils, dataset, model
    alphabets_path = "../datasets/alphabets.txt"
    encoder_pth = "./pths/encoder_150.pth"
    decoder_pth = "./pths/decoder_150.pth"
    input_img_path = "../datasets/CRNN_dataset/test/img/"
    input_gt_path = "../datasets/CRNN_dataset/test/gt/"
    output_path = "../datasets/output/"
    recognize_unit = seq2seq_recognize_unit(alphabets_path,encoder_pth,decoder_pth)
    names =[ i[:-4] for i in os.listdir(input_img_path)]
    for name in tqdm(names):
        img=Image.open(input_img_path+name+".jpg")
        img_tensor = recognize_unit.image_to_tensor(img)
        words = recognize_unit.recognize(img_tensor)
        with open(input_gt_path+name+".txt","r",encoding="utf8") as f:
            box = f.read()
        gt = box+','+words
        #print(gt)
        with open(output_path+name+".txt","w",encoding="utf8") as f:
            f.write(gt)




#print('predict_string: {} => predict_probility: {}'.format(words, prob.item())) #print(gt)