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