File size: 4,595 Bytes
618450b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | """Standalone quantized line recognition; the page pipeline is a separate layer."""
import argparse,collections,hashlib,json
from pathlib import Path
import numpy as np
from rtl import flip
class Recognizer:
def __init__(self,root=None,role='body',provider='CPUExecutionProvider',cache_entries=512,variant='block16'):
self.root=Path(root or Path(__file__).resolve().parent);self.role=role
config=json.loads((self.root/'quantization_config.json').read_text());self.backend=config['format'];self.variant=config.get('default_variant','block16')
self.alphabet=json.loads((self.root/'vocab.json').read_text())['chars'];d=self.root/'components'/role
self.cache=collections.OrderedDict();self.cache_entries=cache_entries;self.calls=0;self.cache_hits=0
if self.backend=='onnx':
import onnxruntime as ort
if provider not in ort.get_available_providers():raise RuntimeError('Unavailable provider: '+provider)
options=ort.SessionOptions();options.intra_op_num_threads=2;options.inter_op_num_threads=1
if variant=='hqq32'and provider!='CPUExecutionProvider':raise ValueError('HQQ32 helper is qualified for ORT CPU only')
name='recognizer-hqq4-block32.onnx'if variant=='hqq32'else'recognizer-w4-block16.onnx'
self.session=ort.InferenceSession(str(d/'models'/name),sess_options=options,providers=[provider]);self.run=lambda x,w:self.session.run(None,dict(pixels=x,widths=w))
elif self.backend=='coreml':
import coremltools as ct
self.session=ct.models.MLModel(str(d/'models/recognizer-w4-block16.mlpackage'),compute_units=ct.ComputeUnit.CPU_AND_GPU)
def run(x,w):
o=self.session.predict(dict(pixels=x,widths=w.astype(np.int32)));return o['tokens'],o['confidence'],o['lengths']
self.run=run
elif self.backend=='litert':
from ai_edge_litert.interpreter import Interpreter
self.session=Interpreter(model_path=str(d/'litert/recognizer-dynamic4-block32.tflite'),num_threads=2);self.runners={}
def run(x,w):
width=x.shape[-1]
if width not in self.runners:self.runners[width]=self.session.get_signature_runner('width'+str(width))
o=self.runners[width](args_0=x,args_1=w.astype(np.int32));return o['output_0'],o['output_1'],o['output_2']
self.run=run
else:raise ValueError(self.backend)
def predict_ink(self,ink):
ink=np.asarray(ink,dtype=np.float32)
if ink.ndim!=2 or ink.shape[0]!=64 or ink.shape[1]<1 or not np.isfinite(ink).all():raise ValueError('Finite float32 ink map [64,true_width] required')
key=hashlib.sha256(str(ink.shape).encode()+ink.tobytes()).hexdigest()
if key in self.cache:
self.cache_hits+=1;self.cache.move_to_end(key);return dict(self.cache[key])
width=ink.shape[1];bucket=max(512,2**int(np.ceil(np.log2(width))))
if bucket>8192:raise ValueError('Width exceeds the validated8192 bound; never truncate a line')
x=np.zeros((1,1,64,bucket),np.float32);x[0,0,:,:width]=ink
tokens,confidence,lengths=self.run(x,np.array([width],np.int64));self.calls+=1
n=int(lengths[0]);seq=np.asarray(tokens[0,:n]);keep=(seq!=0)&np.r_[True,seq[1:]!=seq[:-1]];ids=seq[keep].tolist()
visual=''.join(self.alphabet[int(i)-1]for i in ids)
result=dict(text=flip(visual,base_rtl=True),visual_text=visual,confidence=float(np.asarray(confidence[0,:n])[keep].mean())if keep.any()else 0.,role=self.role)
if self.cache_entries>0:
self.cache[key]=result
while len(self.cache)>self.cache_entries:self.cache.popitem(last=False)
return dict(result)
def read_line(self,image,box=None):
from geometry import crop_line64
gray=np.asarray(image.convert('L'));box=box or [0,0,gray.shape[1],gray.shape[0]]
ink,receipt=crop_line64(gray,box);result=self.predict_ink(ink);result['geometry']=receipt;return result
if __name__=='__main__':
p=argparse.ArgumentParser(description=__doc__);p.add_argument('line_image');p.add_argument('--role',default='body',choices=['body','short_labels','printed','numeric_math','degraded']);p.add_argument('--provider',default='CPUExecutionProvider');p.add_argument('--variant',default='block16',choices=['block16','hqq32']);a=p.parse_args()
from PIL import Image
reader=Recognizer(role=a.role,provider=a.provider,variant=a.variant);print(json.dumps(reader.read_line(Image.open(a.line_image)),ensure_ascii=False,indent=2))
|