File size: 3,384 Bytes
5f70891 | 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 | """Process-local PaddleX adapter for the shared two-branch research recognizer."""
import math
from pathlib import Path
import numpy as np
from PIL import Image
from paddlex.inference.models.predictors.local_model_predictor import LocalModelPredictor
from paddlex.inference.models.text_recognition.predictor import ImageBatchSampler,TextRecResult
from paddlex.inference.models.bindings.registry import register_predictor_binding_map
from shared_character_native import SharedCharacterNative,preprocess_rgb,sha
MODEL_NAME='TrOCR_shared_character_Paddle_research_rec'
class PaddleXSharedCharacter(LocalModelPredictor):
def __init__(self,*args,research_only_allow_unqualified=False,**kwargs):
super().__init__(*args,**kwargs)
if self.model_name!=MODEL_NAME or self.device!='cpu':raise ValueError('Require the shared CPU research model')
native=self.config['Native']
if native['format']!='shared-character-paddle-research-v1' or native['parameters']!=74842752:
raise ValueError('Unexpected shared recognizer profile')
for name,digest in native['files_sha256'].items():
relative=Path(name)
if relative.is_absolute() or '..' in relative.parts or relative.as_posix()!=name:raise ValueError('Invalid bundle asset path')
if sha(self.model_dir/relative)!=digest:raise ValueError('Shared bundle asset bytes changed')
self.runtime=SharedCharacterNative(self.model_dir/'candidate',self.model_dir/'encoder',self.model_dir/'decoders',
allow_unqualified=research_only_allow_unqualified,library_path=self.model_dir/'encoder-library.so')
def _build_batch_sampler(self):return ImageBatchSampler()
def _get_result_class(self):return TextRecResult
def process(self,batch_data,branch=None,return_word_box=False):
if branch not in ['english','character']:raise ValueError('Explicit research branch is required')
if return_word_box:raise ValueError('Autoregressive word-box alignment is not implemented')
images=[];outputs=[]
for item in batch_data.instances:
if isinstance(item,(str,Path)):
with Image.open(item) as source:image=source.convert('RGB')
bgr=np.asarray(image)[:,:,::-1].copy()
elif isinstance(item,np.ndarray) and item.dtype==np.uint8 and item.ndim==3 and item.shape[2]==3:
bgr=item;image=Image.fromarray(np.ascontiguousarray(bgr[:,:,::-1]))
else:raise ValueError('Require an image path or uint8 HWC BGR array')
images.append(bgr);outputs.append(self.runtime.predict_pixels(preprocess_rgb(image),branch))
return dict(input_path=batch_data.input_paths,page_index=batch_data.page_indexes,input_img=images,
rec_text=[r['text'] for r in outputs],rec_score=[min(1.,math.exp(r['score'])) for r in outputs],
beam_score=[r['score'] for r in outputs],tokens=[r['tokens'] for r in outputs],
decoder_iterations=[r['decoder_iterations'] for r in outputs],
token_limit_reached=[r['token_limit_reached'] for r in outputs],recognizer_branch=[branch]*len(outputs),
score_kind=['Uncalibrated geometric mean of generated token probabilities']*len(outputs),vis_font=[None]*len(outputs))
def register():register_predictor_binding_map(PaddleXSharedCharacter,{'flexible':[MODEL_NAME]})
|