| """Local ONNX OCR with shared detection and line assembly for CPU/GPU.""" |
| import base64 |
| import io |
| import json |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import cv2 |
| from PIL import Image, ImageOps |
| from tokenizers import Tokenizer |
|
|
| from .dbnet_utils import DBPostProcess, DetResizeForTest |
| from .line_merger import merge_line_boxes |
|
|
|
|
| def pixels(image, config): |
| size=config['size'] |
| image=image.convert('RGB').resize((size['width'],size['height']),Image.Resampling(config['resample'])) |
| x=np.asarray(image,dtype=np.float32)*np.float32(config['rescale_factor']) |
| x=(x-np.asarray(config['image_mean'],dtype=np.float32))/np.asarray(config['image_std'],dtype=np.float32) |
| return np.ascontiguousarray(x.transpose(2,0,1)[None]) |
|
|
|
|
| def split_line(image, ratio=20): |
| w,h=image.size |
| if w/max(h,1)<=ratio:return [image] |
| ink=(np.asarray(image.convert('L'))<210).sum(axis=0) |
| runs=[];start=None |
| for x,value in enumerate(ink): |
| if value==0 and start is None:start=x |
| elif value!=0 and start is not None:runs.append((start,x));start=None |
| if start is not None:runs.append((start,w)) |
| gaps=[(a,b) for a,b in runs if a>40 and b<w-40 and b-a>=8] |
| if gaps: |
| a,b=min(gaps,key=lambda r:abs((r[0]+r[1])/2-w/2));cut=(a+b)//2 |
| else: |
| left,right=int(w*.35),int(w*.65) |
| cut=left+int(np.argmin(ink[left:right])) if right>left else w//2 |
| cut=max(1,min(w-1,cut)) |
| return split_line(image.crop((0,0,cut,h)),ratio)+split_line(image.crop((cut,0,w,h)),ratio) |
|
|
|
|
| def preview(image): |
| image=image.copy();image.thumbnail((1100,180)) |
| out=io.BytesIO();image.save(out,format='PNG') |
| return 'data:image/png;base64,'+base64.b64encode(out.getvalue()).decode('ascii') |
|
|
|
|
| class OCR: |
| def __init__(self, root=None, backend='cpu', threads=4): |
| import onnxruntime as ort |
| ort.disable_telemetry_events() |
| self.root=Path(root or Path(__file__).resolve().parents[1]) |
| self.config=json.loads((self.root/'config.json').read_text(encoding='utf-8')) |
| self.preprocess=json.loads((self.root/'preprocessor_config.json').read_text(encoding='utf-8')) |
| self.tokenizer=Tokenizer.from_file(str(self.root/'tokenizer.json')) |
| self.tokenizer.no_padding();self.tokenizer.no_truncation() |
| requested_backend=backend |
| if backend not in ('cpu','gpu','auto'):raise ValueError('backend must be cpu, gpu or auto') |
| if backend in ('gpu','auto') and hasattr(ort,'preload_dlls'): |
| ort.preload_dlls() |
| if backend=='auto':backend='gpu' if 'CUDAExecutionProvider' in ort.get_available_providers() else 'cpu' |
| if backend=='gpu' and 'CUDAExecutionProvider' not in ort.get_available_providers(): |
| raise RuntimeError('CUDA недоступна. Установите GPU-зависимости или выберите CPU.') |
| options=ort.SessionOptions();options.intra_op_num_threads=threads;options.inter_op_num_threads=1 |
| if backend=='gpu': |
| variant=self.config['variants']['fp16'] |
| providers=['CUDAExecutionProvider','CPUExecutionProvider'] |
| try: |
| self.encoder=ort.InferenceSession(str(self.root/variant['encoder']),sess_options=options,providers=providers) |
| self.decoder=ort.InferenceSession(str(self.root/variant['decoder']),sess_options=options,providers=providers) |
| if any('CUDAExecutionProvider' not in s.get_providers() for s in (self.encoder,self.decoder)): |
| raise RuntimeError('Модель не загрузилась на CUDA. Выберите CPU; скрытая смена GPU на CPU запрещена.') |
| self.encoder.disable_fallback() |
| self.decoder.disable_fallback() |
| except Exception: |
| if requested_backend!='auto':raise |
| backend='cpu' |
| if backend=='cpu': |
| variant=self.config['variants']['int8'] |
| providers=['CPUExecutionProvider'] |
| self.encoder=ort.InferenceSession(str(self.root/variant['encoder']),sess_options=options,providers=providers) |
| self.decoder=ort.InferenceSession(str(self.root/variant['decoder']),sess_options=options,providers=providers) |
| self.backend=backend |
| self.detector=ort.InferenceSession(str(self.root/'models/dbnet.onnx'),sess_options=options,providers=['CPUExecutionProvider']) |
| self.resize=DetResizeForTest(limit_side_len=736,limit_type='min') |
| self.post=DBPostProcess(thresh=.2,box_thresh=.2,max_candidates=1000,unclip_ratio=1.3,use_dilation=True,score_mode='fast') |
|
|
| def recognize_line(self,image): |
| hidden=self.encoder.run(None,{'pixel_values':pixels(image,self.preprocess)})[0] |
| ids=np.array([[self.config['decoder_start_token_id']]],dtype=np.int64) |
| truncated=True |
| for _ in range(self.config['max_new_tokens']): |
| logits=self.decoder.run(None,{'input_ids':ids,'encoder_hidden_states':hidden})[0] |
| if not np.isfinite(logits).all():raise RuntimeError('Non-finite decoder output') |
| token=int(np.argmax(logits[0,-1])) |
| ids=np.concatenate((ids,np.array([[token]],dtype=np.int64)),axis=1) |
| if token==self.config['eos_token_id']:truncated=False;break |
| return self.tokenizer.decode(ids[0].tolist(),skip_special_tokens=True).strip(),truncated |
|
|
| def detect(self,image): |
| |
| |
| scale=max(1.,736/min(image.size)) |
| if max(image.size)*scale>4096: |
| scale=4096/max(image.size) |
| width=max(32,round(image.width*scale/32)*32) |
| height=max(32,round(image.height*scale/32)*32) |
| data=dict(image=cv2.resize(np.asarray(image),(width,height)), |
| shape=np.array([image.height,image.width,height/image.height,width/image.width])) |
| else:data=self.resize({'image':np.asarray(image)}) |
| x=np.asarray(data['image'],dtype=np.float32)*np.float32(1/255) |
| x=(x-np.array([.485,.456,.406],dtype=np.float32))/np.array([.229,.224,.225],dtype=np.float32) |
| x=np.ascontiguousarray(x.transpose(2,0,1)[None]) |
| output=self.detector.run(None,{self.detector.get_inputs()[0].name:x})[0] |
| boxes=self.post(output,np.asarray([data['shape']]))[0]['points'] |
| valid=[] |
| for box in boxes: |
| box=np.asarray(box,dtype=np.float32) |
| box[:,0]=np.clip(box[:,0],0,image.width-1);box[:,1]=np.clip(box[:,1],0,image.height-1) |
| if np.linalg.norm(box[0]-box[1])>3 and np.linalg.norm(box[0]-box[3])>3:valid.append(box) |
| return merge_line_boxes(valid,x_margin_ratio=3.5,y_overlap_thresh=.45) |
|
|
| def recognize(self,image,boxes=None,include_images=True): |
| t=time.perf_counter() |
| image=ImageOps.exif_transpose(image).convert('RGB') |
| if np.asarray(image.convert('L')).mean()<125:image=ImageOps.invert(image) |
| boxes=self.detect(image) if boxes is None else boxes |
| lines=[] |
| for box in boxes: |
| a,b,c,d=[int(v) for v in box] |
| if c-a<4 or d-b<4:continue |
| crop=image.crop((max(0,a-4),max(0,b-2),min(image.width,c+4),min(image.height,d+2))) |
| parts=[self.recognize_line(ch) for ch in split_line(crop)] |
| text=' '.join(s for s,_ in parts if s).strip() |
| line=dict(bbox=[a,b,c,d],text=text,truncated=any(end for _,end in parts),chunks=len(parts)) |
| if include_images:line['image']=preview(crop) |
| lines.append(line) |
| return dict(backend=self.backend,lines=lines,full_text='\n'.join(row['text'] for row in lines), |
| seconds=time.perf_counter()-t,no_text_detected=not lines) |
|
|
