get_n_best_hyps.py

import pickle
import argparse

import torch
import torch.backends.cudnn as cudnn
import torch.utils.data
from tqdm import tqdm

from utils import CTCLabelConverter
from dataset import LmdbDataset, AlignCollate
from model import Model
from tools.ctc_utils import ctc_prefix_beam_search
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def test(opt):
    converter = CTCLabelConverter(opt.character)
    opt.num_class = len(converter.character)

    if opt.rgb:
        opt.input_channel = 3
    model = Model(opt)
    print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,
          opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction,
          opt.SequenceModeling, opt.Prediction)
    model = torch.nn.DataParallel(model).to(device)

    # load model
    print('loading pretrained model from %s' % opt.saved_model)
    model.load_state_dict(torch.load(opt.saved_model, map_location=device))

    dataset = LmdbDataset(opt.eval_data, opt)
    AlignCollate_evaluation = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)

    model.eval()
    with torch.no_grad():
        cache = {}
        for idx in tqdm(range(opt.start_index, len(dataset)), desc="Processing", ncols=100):
            sample = dataset[idx]
            image_tensors, _, file_names, _ = AlignCollate_evaluation([sample])
            batch_size = image_tensors.size(0)
            image = image_tensors.to(device)
            text_for_pred = torch.LongTensor(batch_size, opt.batch_max_length + 1).fill_(0).to(device)
        
            preds = model(image, text_for_pred)
            hyps, scores = ctc_prefix_beam_search(preds, beam_size=opt.beam_size)

            n_best_strs = ["".join([converter.idict[i] for i in hyp[0]]) for hyp in hyps]
            cache[('cor-%09d' % int(file_names[0].split('-')[-1])).encode()] = pickle.dumps({'str': n_best_strs, 'conf': scores})

            if len(cache) % 1000 == 0:
                dataset.writeCache(cache)
                cache = {}

        dataset.writeCache(cache)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--eval_data', required=True, help='path to evaluation dataset')
    parser.add_argument('--workers', type=int, help='number of data loading workers', default=4)
    parser.add_argument('--batch_size', type=int, default=192, help='input batch size')
    parser.add_argument('--beam_size', type=int, default=1, help='ctc beam search size')
    parser.add_argument('--saved_model', required=True, help="path to saved_model to evaluation")
    """ Data processing """
    parser.add_argument('--batch_max_length', type=int, default=25, help='maximum-label-length')
    parser.add_argument('--imgH', type=int, default=32, help='the height of the input image')
    parser.add_argument('--imgW', type=int, default=100, help='the width of the input image')
    parser.add_argument('--rgb', action='store_true', help='use rgb input')
    parser.add_argument('--character', type=str, default='0123456789abcdefghijklmnopqrstuvwxyz', help='character label')
    parser.add_argument('--sensitive', action='store_true', help='for sensitive character mode')
    parser.add_argument('--PAD', action='store_true', help='whether to keep ratio then pad for image resize')
    parser.add_argument('--data_filtering_off', action='store_true', help='for data_filtering_off mode')
    parser.add_argument('--baiduCTC', action='store_true', help='for data_filtering_off mode')
    """ Model Architecture """
    parser.add_argument('--Transformation', type=str, required=True, help='Transformation stage. None|TPS')
    parser.add_argument('--FeatureExtraction', type=str, required=True, help='FeatureExtraction stage. VGG|RCNN|ResNet')
    parser.add_argument('--SequenceModeling', type=str, required=True, help='SequenceModeling stage. None|BiLSTM')
    parser.add_argument('--Prediction', type=str, required=True, help='Prediction stage. CTC|Attn')
    parser.add_argument('--num_fiducial', type=int, default=20, help='number of fiducial points of TPS-STN')
    parser.add_argument('--input_channel', type=int, default=1, help='the number of input channel of Feature extractor')
    parser.add_argument('--output_channel', type=int, default=512, help='the number of output channel of Feature extractor')
    parser.add_argument('--hidden_size', type=int, default=256, help='the size of the LSTM hidden state')
    """ Calibration Ralated """
    parser.add_argument('--start_index', type=int, default=0, help='the size of the LSTM hidden state')
    parser.add_argument('--with_file_name', action='store_true')
    parser.add_argument('--with_vis', action='store_true')
    opt = parser.parse_args()
    """ vocab / character number configuration """
    if opt.sensitive:
        with open('data_lmdb_release/charset_94.txt','r') as f:
            opt.character = "".join(sorted(f.read()))
    else:
        with open('data_lmdb_release/charset_36.txt','r') as f:
            opt.character = "".join(sorted(f.read()))

    cudnn.benchmark = True
    cudnn.deterministic = True
    opt.num_gpu = torch.cuda.device_count()

    test(opt)