| |
| """Bina 0.4 Rizeh: Persian page OCR (Surya layout + 95k Indra CTC student + printed/table/numeric experts and assembly). |
| |
| usage: python rizeh.py page.jpg [more.png doc.pdf ...] [--out results] [--device cuda|cpu] |
| """ |
| import argparse,json,os,sys,time,importlib.util |
| from pathlib import Path |
| ROOT=Path(__file__).resolve().parent;ST=ROOT/'stallion';MARK=ROOT/'.bina-root' |
| LAYOUT_REPO,LAYOUT_REV='datalab-to/surya_layout2','0aee81d5fd9275c0582e545bf3a56944b1e75679' |
| TEXT=('.py','.json','.yml','.yaml','.txt','.jsonl') |
|
|
| def setup(): |
| """Point every bundled path at this checkout (idempotent; reruns after a move) and fetch the pinned layout model.""" |
| here=str(ROOT);old=MARK.read_text().strip()if MARK.exists()else'__BINA_ROOT__' |
| if old!=here: |
| n=0 |
| for p in ST.rglob('*'): |
| if p.is_file()and p.suffix in TEXT: |
| s=p.read_text(encoding='utf-8') |
| if old in s:p.write_text(s.replace(old,here),encoding='utf-8');n+=1 |
| MARK.write_text(here);print(f'[rizeh] bound {n} files to {here}',file=sys.stderr) |
| snap=ST/'indra-scale/hf/hub/models--datalab-to--surya_layout2/snapshots'/LAYOUT_REV |
| if not(snap/'rfdetr_layout.pth').exists(): |
| from huggingface_hub import snapshot_download |
| print('[rizeh] downloading layout model',LAYOUT_REPO,LAYOUT_REV,file=sys.stderr) |
| snapshot_download(LAYOUT_REPO,revision=LAYOUT_REV,cache_dir=str(ST/'indra-scale/hf/hub')) |
| return snap |
|
|
| def pages(path): |
| p=Path(path) |
| if p.suffix.lower()=='.pdf': |
| import pypdfium2 as pdfium |
| doc=pdfium.PdfDocument(str(p)) |
| for i in range(len(doc)):yield f'{p.stem}-p{i+1:03d}',doc[i].render(scale=200/72).to_pil().convert('RGB') |
| else: |
| from PIL import Image |
| yield p.stem,Image.open(p).convert('RGB') |
|
|
| class Rizeh: |
| def __init__(self,device='cuda'): |
| snap=setup() |
| os.environ.setdefault('GATE_DEVICE',device);os.environ.setdefault('FAST_DETECTOR_DEVICE','cpu');os.environ.setdefault('DETECTOR_IMAGE_CHUNK_HEIGHT','1400');os.environ['FAST_ORDER_MODEL_CHECKPOINT']=str(snap/'order') |
| P=ST/'indra-scale';sys.path[:0]=[str(P),str(ST/'bina04-pagepath-v3-round3-clean-20261031'),str(P/'scoringTT')] |
| spec=importlib.util.spec_from_file_location('bina04_rizeh_runtime',ST/'bina04-ink-loss-guard-95k-20261104/production.py');m=importlib.util.module_from_spec(spec);spec.loader.exec_module(m) |
| self.engine=m.Runtime(device) |
| from surya.fast_layout.server import LayoutEngine |
| self.layout=LayoutEngine(checkpoint=str(snap)) |
| def __call__(self,image,name='page'): |
| raw=self.layout.run_batch([image],[dict(threshold=None,use_order=True)])[0].model_dump(mode='json') |
| out=self.engine(image,dict(raw_layout=raw),'page',name) |
| lines=[dict(text=l.get('text',''),layer=l.get('layer'),box=[round(float(v),1)for v in l['box']],confidence=l.get('confidence'))for l in out.get('lines',[])if l.get('box')] |
| return dict(text=out['text'],lines=lines) |
| def close(self):self.engine.close() |
|
|
| def main(): |
| ap=argparse.ArgumentParser(description=__doc__,formatter_class=argparse.RawDescriptionHelpFormatter);ap.add_argument('inputs',nargs='+');ap.add_argument('--out',default='rizeh-out');ap.add_argument('--device',default='cuda',choices=['cuda','cpu']);a=ap.parse_args() |
| O=Path(a.out);O.mkdir(parents=True,exist_ok=True);r=Rizeh(a.device) |
| for src in a.inputs: |
| for name,im in pages(src): |
| t=time.perf_counter();res=r(im,name);res.update(source=str(src),seconds=round(time.perf_counter()-t,2),model='Bina-0.4-Rizeh') |
| (O/f'{name}.json').write_text(json.dumps(res,ensure_ascii=False,indent=1),encoding='utf-8');(O/f'{name}.txt').write_text(res['text'],encoding='utf-8') |
| print(f'{name}: {len(res["lines"])} lines, {res["seconds"]}s -> {O/name}.txt',flush=True) |
| r.close() |
| if __name__=='__main__':main() |
|
|
