File size: 7,812 Bytes
c8cdee4 | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | """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):
# The original detector upsizes the short side. Bound unusually narrow
# images so a one-line crop cannot allocate a gigantic detector tensor.
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)
|