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