src/text_recognition/text_recognition.py

# Copyright (c) 2022, National Diet Library, Japan
#
# This software is released under the CC BY 4.0.
# https://creativecommons.org/licenses/by/4.0/


import argparse
import functools
import difflib
import collections
import pathlib
from tqdm import tqdm

import xml.etree.ElementTree as ET
from xml.dom import minidom

from PIL import Image, ImageDraw, ImageFont
import torch
import torch.backends.cudnn as cudnn
import torch.nn.functional as F
from torch.utils.data import ConcatDataset, Subset
from nltk.metrics import edit_distance

from utils import CTCLabelConverter, AttnLabelConverter
from dataset import XMLLmdbDataset, AlignCollate, tensor2im
from model import Model
from xmlparser import XMLRawDataset, SyntheticDataset, XMLRawDatasetWithCli
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def gen_dataset(db_type, db_path, opt, line_index=None, accept_empty=True, keep_remain=False):
    if db_type == 'xmllmdb':
        ds = ConcatDataset([XMLLmdbDataset(root=p, opt=opt) for p in db_path])
        if line_index is not None:
            ds = Subset(ds, opt.line_index)
    elif db_type == 'xmlraw':
        ds = XMLRawDataset.from_list(input_paths=db_path,
                                     image_type=XMLRawDataset.IMAGE_TYPE_GRAY_IMAGE,
                                     accept_empty=accept_empty, keep_remain=keep_remain)
        opt.workers = 0
    elif db_type == 'synth':
        ds = SyntheticDataset(opt.character, db_path)
    return ds


def _debug_char_prob(preds_prob, character):
    preds_v, preds_i = torch.topk(preds_prob, 3)
    for b in zip(preds_v.tolist(), preds_i.tolist()):
        for p3, i3 in zip(*b):
            if i3[0] == 0:
                continue
            for p, i in zip(p3, i3):
                if p > 0.01:
                    print(f'{p:.2f}', character[i-1], end=' ')
            print()
        print('--------')


class Inferencer:
    @staticmethod
    def get_argparser():
        parser = argparse.ArgumentParser()
        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('--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('--remove_char', default=None, help='remove the specified index class. ex. 〓')
        """ 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')
        return parser

    def __init__(self, opt):
        """
        Args:
            opt
                上記get_parserによってparseされたargument
        """
        # model config
        if 'CTC' in opt.Prediction:
            converter = CTCLabelConverter(opt.character)
        else:
            converter = AttnLabelConverter(opt.character)
        opt.num_class = len(converter.character)
        if opt.remove_char is not None:
            opt.remove_char = opt.character.index(opt.remove_char) + 1

        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))

        self.model = model
        self.converter = converter
        self.aligncollate = AlignCollate(imgH=opt.imgH, imgW=opt.imgW, keep_ratio_with_pad=opt.PAD)
        self.opt = opt

    def gen(self, dataset, keep_remain=False, with_tqdm=False):
        """
        Args:
            dataset
                以下を生成するtorch.utils.data.Dataset
                PIL.Image(mode="L"), {'WIDTH': int, 'HEIGHT': int, 'STRING': string}
            keep_remain
                これが有効のとき、xmlraw dbは偶数週目に
                推論しない要素を吐くようになる
            with_tqdm
                これが有効のとき、進捗表示をする
        Yields:
            image
            groundtruth label
            prediction label
            confidence score
            appendix information
        """
        converter = self.converter

        demo_loader = torch.utils.data.DataLoader(
            dataset, batch_size=self.opt.batch_size,
            shuffle=False,
            num_workers=int(self.opt.workers),
            collate_fn=self.aligncollate, pin_memory=True)
        if with_tqdm:
            demo_loader = tqdm(demo_loader, ncols=80)

        # predict
        self.model.eval()
        with torch.no_grad():
            for image_tensors, labels, data in demo_loader:
                batch_size = image_tensors.size(0)
                image = image_tensors.to(device)
                # For max length prediction
                length_for_pred = torch.IntTensor([self.opt.batch_max_length] * batch_size).to(device)
                text_for_pred = torch.LongTensor(batch_size, self.opt.batch_max_length + 1).fill_(0).to(device)

                if 'CTC' in self.opt.Prediction:
                    preds = self.model(image, text_for_pred)
                    if self.opt.remove_char is not None:
                        preds[:, :, self.opt.remove_char] = -1e5

                    # Select max probabilty (greedy decoding) then decode index to character
                    preds_size = torch.IntTensor([preds.size(1)] * batch_size)
                    _, preds_index = preds.max(2)
                    # preds_index = preds_index.view(-1)
                    preds_str = converter.decode(preds_index, preds_size)
                else:
                    preds = self.model(image, text_for_pred, is_train=False)

                    # select max probabilty (greedy decoding) then decode index to character
                    _, preds_index = preds.max(2)
                    preds_str = converter.decode(preds_index, length_for_pred)

                preds_prob = F.softmax(preds, dim=2)
                preds_max_prob, _ = preds_prob.max(dim=2)

                if 0:
                    _debug_char_prob(preds_prob, self.opt.character)

                for image, gt, pred, pred_max_prob, datum in zip(image, labels, preds_str, preds_max_prob, data):
                    if 'Attn' in self.opt.Prediction:
                        pred_EOS = pred.find('[s]')
                        pred = pred[:pred_EOS]  # prune after "end of sentence" token ([s])
                        pred_max_prob = pred_max_prob[:pred_EOS]

                    # calculate confidence score (= multiply of pred_max_prob)
                    try:
                        confidence_score = pred_max_prob.cumprod(dim=0)[-1]
                    except Exception:
                        confidence_score = 0  # for empty pred case, when prune after "end of sentence" token ([s])

                    yield image, gt, pred, confidence_score, datum

            if keep_remain:
                for datum in dataset:
                    yield None, None, None, None, datum


class TR_WORKER:
    CHAR_DIFF_NONE = 0
    CHAR_DIFF_WRONG = 1
    CHAR_DIFF_ALL = 2

    def __init__(self,
                 accuracy=False,
                 levenshtein_distance=False,
                 char_diff=CHAR_DIFF_NONE,
                 render=False,
                 xml=None,
                 outimage_dir=None, font_path=None,
                 stat=False):
        self._task = []
        self._accuracy = accuracy
        self._char_diff = char_diff
        self._levenshtein_distance = levenshtein_distance
        self._xml = xml
        self._stat = stat

        self.nline = 0
        if accuracy:
            self.accuracy = 0
            self.ncorrect = 0
            self._task.append(self._facc)

        if levenshtein_distance:
            self.sum_dist = 0
            self.normalized_edit_distance = 0
            self._task.append(self._fld)

        if char_diff != self.CHAR_DIFF_NONE:
            self.counters = {
                'misstake': collections.Counter(),
                'tp': collections.Counter(),
                'fn': collections.Counter(),
                'fp': collections.Counter(),
            }
            self._task.append(self._fchar_diff)

            self.outimage_dir = outimage_dir
            if outimage_dir is None:
                self.outimage_dir = None
            else:
                assert font_path is not None
                self.outimage_dir = pathlib.Path(outimage_dir)
                self.outimage_dir.mkdir(exist_ok=True)
                dtmp = ImageDraw.Draw(Image.new('L', (400, 200)))
                self._font = ImageFont.truetype(font_path, 32)
                self._textsize = functools.partial(dtmp.multiline_textsize, font=self._font)

        if render:
            self._task.append(self._frender)
            assert font_path is not None
            assert outimage_dir is not None
            self.outimage_dir = pathlib.Path(outimage_dir)
            self.outimage_dir.mkdir(exist_ok=True)
            dtmp = ImageDraw.Draw(Image.new('L', (400, 200)))
            self._font = ImageFont.truetype(font_path, 32)
            self._textsize = functools.partial(dtmp.multiline_textsize, font=self._font)

        if xml:
            self.outxml_dir = pathlib.Path(xml)
            self.outxml_dir.mkdir(exist_ok=True)
            self._xml_data = {}
            self._task.append(self._fxml)

    def finalize(self):
        if self._accuracy:
            self.accuracy = self.ncorrect / self.nline
        if self._levenshtein_distance:
            self.normalized_edit_distance = self.sum_dist / self.nline
        if self._xml:
            self._fgenerate_xml()
        if self._stat:
            print('===== f measure =====')
            for c in self.counters['tp'].keys() | self.counters['fp'].keys() | self.counters['fn'].keys():
                tp = self.counters['tp'][c]
                precision = tp / (tp + self.counters['fp'][c] + 1e-9)
                recall = tp / (tp + self.counters['fn'][c] + 1e-9)
                print(c, f"{2 * precision * recall / (precision + recall + 1e-9):.3f}")
            print('===== misstake stat =====')
            for p, n in self.counters['misstake'].most_common():
                if p[1] == '-':
                    print(p, n, f"U+{ord(p[0]):X} U+{ord(p[2]):X}")
        return self

    def _facc(self, correct, *args):
        if correct:
            self.ncorrect += 1

    def _fld(self, correct, image, gt, pred, *args):
        d = edit_distance(gt, pred)
        if len(gt) == 0 and len(pred) == 0:
            self.sum_dist += 0
        elif len(gt) > len(pred):
            self.sum_dist += 1 - d / len(gt)
        else:
            self.sum_dist += 1 - d / len(pred)

    def _frender(self, correct, image, sa1, sb1, *args):
        image_pil = Image.fromarray(tensor2im(image))
        w, h = self._textsize(f'{sb1}')
        g = Image.new(image_pil.mode, (w, h), (255, 255, 255))
        d = ImageDraw.Draw(g)
        p = [0, 0]
        draw_escape_colored_text(sb1, d, p=p, font=self._font)
        if h * image_pil.width > image_pil.height * 2 * w:
            w = w * image_pil.height * 2 // h
            h = image_pil.height * 2
        else:
            h = h * image_pil.width // w
            w = image_pil.width
        g = g.resize((w, h))
        canvas = Image.new(image_pil.mode, (image_pil.width, image_pil.height + h), (255, 255, 255))
        canvas.paste(image_pil)
        canvas.paste(g, (0, image_pil.height))
        canvas.save(self.outimage_dir / f'{self.nline:09d}-{sb1.replace("/", "")}.png')

    def _fchar_diff(self, correct, image, sa1, sb1, *args):
        if correct and self._char_diff != self.CHAR_DIFF_ALL:
            if self._char_diff == self.CHAR_DIFF_ALL:
                print('------------------')
                print(sa1)
            return
        if sa1 is None:
            sa1 = ''
        sm = difflib.SequenceMatcher(None, sa1, sb1)
        sa2 = str()
        sb2 = str()
        reason = ''
        for tag, ia1, ia2, ib1, ib2 in sm.get_opcodes():
            if tag == 'equal':
                sa2 += "\033[0m"
                sb2 += "\033[0m"
                self.counters['tp'].update(list(sa1[ia1:ia2]))
            elif tag == 'replace':
                sa2 += "\033[31m"
                sb2 += "\033[31m"
                self.counters['fn'].update(list(sa1[ia1:ia2]))
                self.counters['fp'].update(list(sb1[ia1:ib2]))
                for ia, ib in zip(range(ia1, ia2), range(ib1, ib2)):
                    self.counters['misstake'].update([f'{sa1[ia]}-{sb1[ib]}'])
                    reason += f'{sa1[ia]}-{sb1[ib]},'
            elif tag == 'insert':
                sb2 += "\033[33m"
                self.counters['fp'].update(list(sb1[ia1:ib2]))
                for ia in range(ia1, ia2):
                    self.counters['misstake'].update([f'{sa1[ia]}>　'])
                    reason += f'{sa1[ia]}>　,'
            elif tag == 'delete':
                sa2 += "\033[33m"
                self.counters['fn'].update(list(sa1[ia1:ia2]))
                for ib in range(ib1, ib2):
                    self.counters['misstake'].update([f'　<{sb1[ib]}'])
                    reason += f'　<{sb1[ib]},'
            sa2 += sa1[ia1:ia2]
            sb2 += sb1[ib1:ib2]
        sa2 += '\033[0m'
        sb2 += '\033[0m'

        if self._char_diff != self.CHAR_DIFF_NONE:
            print(f'-{self.nline:09d}-----------------')
            print(sa2)
            print(sb2)

        if self.outimage_dir is not None:
            image_pil = Image.fromarray(tensor2im(image))
            w, h = self._textsize(f'{sa2}\n{sb2}')
            g = Image.new(image_pil.mode, (w, h), (255, 255, 255))
            d = ImageDraw.Draw(g)
            p = [0, 0]
            draw_escape_colored_text(sa2, d, p=p, font=self._font)
            draw_escape_colored_text(sb2, d, p=p, font=self._font)
            if h * image_pil.width > image_pil.height * 4 * w:
                w = w * image_pil.height * 4 // h
                h = image_pil.height * 4
            else:
                h = h * image_pil.width // w
                w = image_pil.width
            g = g.resize((w, h))
            canvas = Image.new(image_pil.mode, (image_pil.width, image_pil.height + h), (255, 255, 255))
            canvas.paste(image_pil)
            canvas.paste(g, (0, image_pil.height))
            # canvas.save(self.outimage_dir / f'{self.nline:09d}-{reason}.png')
            canvas.save(self.outimage_dir / f'{self.nline:09d}-{sa1.replace("/", "")}.png')

    def _fxml(self, _1, _2, _3, pred_str, conf, data):
        d = dict()
        for attr in ['tag', 'DIRECTION', 'TYPE', 'X', 'Y', 'WIDTH', 'HEIGHT', 'CONF']:
            if attr in data:
                d[attr] = f"{data[attr]}"
        if conf is not None:
            d['STR_CONF'] = f"{conf:.3f}"
        if pred_str is not None:
            d['STRING'] = pred_str
        pid = data['path'].parents[1].name
        imagename = data['path'].name

        if pid not in self._xml_data:
            self._xml_data[pid] = {}
        if imagename not in self._xml_data[pid]:
            self._xml_data[pid][imagename] = []

        self._xml_data[pid][imagename].append(d)

    def _fgenerate_xml(self):
        for pid, pages in self._xml_data.items():
            xml_data = ET.Element('OCRDATASET')
            ET.register_namespace('', 'NDLOCRDATASET')
            for p, lines in pages.items():
                page = ET.SubElement(xml_data, 'PAGE', attrib={'IMAGENAME': p})
                for line in lines:
                    line = ET.SubElement(page, line.pop('tag', 'LINE'), attrib=line)

            xml_str = minidom.parseString(ET.tostring(xml_data, encoding='utf-8', method='xml')).toprettyxml(indent=' ')
            out_xml_path = self.outxml_dir / (pid + '.xml')

            with out_xml_path.open(mode='w') as f:
                f.write(xml_str)

    def __call__(self, generator):
        for image, gt, pred, conf, data in generator:
            correct = gt == pred
            for t in self._task:
                t(correct, image, gt, pred, conf, data)
            self.nline += 1
        return self


def draw_escape_colored_text(t, d, p, font):
    get_textsize = functools.partial(d.textsize, font=font)
    it = iter(t)
    cl = (0, 0, 0)
    for c in it:
        if c == '\033':
            n = next(it)
            while n[-1] != 'm':
                n += next(it)
            if n == '[0m':
                cl = (0, 0, 0)
            elif n == '[31m':
                cl = (255, 0, 0)
            elif n == '[33m':
                cl = (255, 255, 0)
            continue
        else:
            size = get_textsize(c)

        d.text(p, c, font=font, fill=cl)
        p[0] += size[0]
    p[0], p[1] = 0, get_textsize(t)[1]


class InferencerWithCLI:
    def __init__(self, conf_dict, character):
        class EmptyOption():
            def __init__(self):
                return

        # create option dictionary from parser
        parser = Inferencer.get_argparser()
        option_key_dict = {}
        for action in parser._actions:
            for opt_str in action.option_strings:
                key_str = None
                if opt_str.startswith('--'):
                    key_str = opt_str[2:]
                    option_key_dict[key_str] = parser.get_default(key_str)

        # create option instance
        opt = EmptyOption()
        for k, v in option_key_dict.items():
            setattr(opt, k, v)
        opt.saved_model = conf_dict['saved_model']
        opt.batch_max_length = conf_dict['batch_max_length']
        opt.batch_size = conf_dict['batch_size']
        opt.character = character
        opt.imgW = conf_dict['imgW']
        opt.workers = conf_dict['workers']
        opt.xml = conf_dict['xml']
        opt.FeatureExtraction = conf_dict['FeatureExtraction']
        opt.Prediction = conf_dict['Prediction']
        opt.PAD = conf_dict['PAD']
        opt.SequenceModeling = conf_dict['SequenceModeling']
        opt.Transformation = conf_dict['Transformation']

        self.opt = opt
        self.inf = Inferencer(self.opt)

        return

    def inference_wich_cli(self, img_data, xml_data, accept_empty=False,
                           yield_block_pillar=True, yield_block_page_num=True):

        cudnn.benchmark = True
        cudnn.deterministic = True
        num_gpu = torch.cuda.device_count()
        dataset = XMLRawDatasetWithCli(img_data, xml_data,
                                       accept_empty=accept_empty,
                                       yield_block_pillar=yield_block_pillar,
                                       yield_block_page_num=yield_block_page_num)
        generator = self.inf.gen(dataset, keep_remain=self.opt.xml)

        result_list = []
        for image, gt, pred, conf, data in generator:
            result_list.append(pred)

        for xml_line in xml_data.getroot().find('PAGE'):
            if len(result_list) == 0:
                print('ERROR: mismatch num of predicted result and xml line')
                break
            if result_list[0] is None:
                print('No predicted STRING for this xml_line')
                print(xml_line.attrib)
                del result_list[0]
                continue
            xml_line.set('STRING', result_list.pop(0))

        return xml_data


if __name__ == '__main__':
    parser = Inferencer.get_argparser()
    g = parser.add_argument_group('db settings')
    g.add_argument('--db_path', required=True, nargs='+', help='データベースへのパス(複数指定可). synthの場合はfont pathを指定する')
    g.add_argument('--db_type', choices=['xmlraw', 'xmllmdb', 'synth'], help='データベースの種類', default='xmlraw')
    g.add_argument('--line_index', type=int, nargs='+', default=None, help='指定の行のみに対して推論. xmllmdb使用時のみ有効')
    action = parser.add_argument_group()
    action.add_argument('--diff', nargs='?', default='none', const='wrong', choices=['none', 'wrong', 'all'],
                        help='差分表示. 画像出力したい場合にはoutimage_dirとfont_pathを指定する')
    action.add_argument('--render', action='store_true', help='diffのgtなし番. outimage_dirとfont_pathが必要')
    action.add_argument('--leven', action='store_true', help='normalized edit distance')
    action.add_argument('--acc', action='store_true', help='accuracy')
    action.add_argument('--xml', default=None, help='xml出力を行う先を指定する')
    parser.add_argument('--stat', action='store_true', help='diff指定時の出力を詳細にする')
    parser.add_argument('--outimage_dir', default=None, help='diff指定時の画像出力先')
    parser.add_argument('--font_path', default=None, help='diff指定時画像出力する際に使用するttf font')
    parser.add_argument('--skip_empty', dest='accept_empty', action='store_false', help='GTが空行の推論を行わない')
    opt = parser.parse_args()

    assert opt.diff != 'none' or opt.render or opt.leven or opt.acc or opt.xml

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

    dataset = gen_dataset(opt.db_type, opt.db_path, opt, line_index=opt.line_index,
                          accept_empty=opt.accept_empty, keep_remain=opt.xml)
    generator = Inferencer(opt).gen(dataset, keep_remain=opt.xml, with_tqdm=True)
    char_diff = {
        'none': TR_WORKER.CHAR_DIFF_NONE,
        'wrong': TR_WORKER.CHAR_DIFF_WRONG,
        'all': TR_WORKER.CHAR_DIFF_ALL,
    }[opt.diff]
    w = TR_WORKER(char_diff=char_diff, render=opt.render, stat=opt.stat,
                  accuracy=opt.acc, levenshtein_distance=opt.leven,
                  xml=opt.xml,
                  outimage_dir=opt.outimage_dir,
                  font_path=opt.font_path)(generator).finalize()
    if w._accuracy:
        print(f'Accuracy: {w.accuracy:.4f}')
    if w._levenshtein_distance:
        print(f'Normalized Edit Distance: {w.normalized_edit_distance:.4f}')