Spaces:
Running on Zero
Running on Zero
| """CPU LayoutLMv3 extraction in an isolated legacy-Transformers environment.""" | |
| import json | |
| import os | |
| import sys | |
| from pathlib import Path | |
| def run(pdf, destination): | |
| import fitz | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from huggingface_hub import hf_hub_download | |
| from transformers import modeling_utils, pytorch_utils | |
| for name in ('find_pruneable_heads_and_indices', 'prune_linear_layer'): | |
| if not hasattr(modeling_utils, name): | |
| setattr(modeling_utils, name, getattr(pytorch_utils, name)) | |
| from detectron2.config import get_cfg | |
| from detectron2.engine import DefaultPredictor | |
| from layout_vendor.model_init import add_vit_config | |
| import layout_vendor.backbone | |
| import layout_vendor.rcnn_vl | |
| torch.set_num_threads(max(1, min(4, os.cpu_count() or 2))) | |
| revision = 'ed6b654c018d742e65a17671e379c5e6ecc87ec9' | |
| model_dir = None | |
| for name in ('config.json', 'model_final.pth'): | |
| p = hf_hub_download('opendatalab/PDF-Extract-Kit-1.0', | |
| 'models/Layout/LayoutLMv3/' + name, revision=revision, token=False) | |
| model_dir = str(Path(p).parent) | |
| cfg = get_cfg() | |
| add_vit_config(cfg) | |
| cfg.merge_from_file(str(Path(__file__).parent/'layout_vendor/layoutlmv3_base_inference.yaml')) | |
| cfg.MODEL.DEVICE = 'cpu' | |
| cfg.MODEL.CONFIG_PATH = model_dir | |
| cfg.MODEL.WEIGHTS = str(Path(model_dir)/'model_final.pth') | |
| cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 | |
| cfg.freeze() | |
| predictor = DefaultPredictor(cfg) | |
| classes = ['title', 'plain text', 'abandon', 'figure', 'figure_caption', 'table', | |
| 'table_caption', 'table_footnote', 'isolate_formula', 'formula_caption'] | |
| kinds = {'figure':'figure','table':'table','isolate_formula':'formula'} | |
| out = Path(destination); (out/'assets').mkdir(parents=True, exist_ok=True) | |
| layout, assets, counters = [], [], {k:0 for k in kinds} | |
| def iou(a, b): | |
| w = max(0, min(a[2],b[2])-max(a[0],b[0])) | |
| h = max(0, min(a[3],b[3])-max(a[1],b[1])) | |
| return w*h/max(1, (a[2]-a[0])*(a[3]-a[1])+(b[2]-b[0])*(b[3]-b[1])-w*h) | |
| with fitz.open(pdf) as doc: | |
| if doc.needs_pass: | |
| raise ValueError('Password-protected PDFs are not supported.') | |
| n = min(12, len(doc)) | |
| for i in range(n): | |
| page = doc[i] | |
| if page.rect.width*page.rect.height*(200/72)**2 > 20_000_000: | |
| raise ValueError('PDF page dimensions are too large.') | |
| pix = page.get_pixmap(matrix=fitz.Matrix(200/72,200/72), colorspace=fitz.csRGB, alpha=False) | |
| image = Image.frombytes('RGB', (pix.width,pix.height), pix.samples) | |
| pred = predictor(np.asarray(image)[:,:,::-1])['instances'].to('cpu') | |
| dets=[] | |
| for box, score, cls in zip(pred.pred_boxes.tensor.tolist(),pred.scores.tolist(),pred.pred_classes.tolist()): | |
| dets.append({'page':i+1,'class_name':classes[cls], 'bbox':box,'score':score}) | |
| layout.extend(dets) | |
| selected=[] | |
| for d in sorted(dets,key=lambda d:-d['score']): | |
| b=d['bbox']; area=(b[2]-b[0])*(b[3]-b[1]) | |
| if d['class_name'] not in kinds or d['score']<.60 or area<.004*pix.width*pix.height: | |
| continue | |
| if any(iou(b,k['bbox'])>=.75 for k in selected): | |
| continue | |
| selected.append(d) | |
| cls=d['class_name']; counters[cls]+=1; aid=f'{kinds[cls]}_{counters[cls]}' | |
| crop=image.crop((max(0,int(b[0])-6),max(0,int(b[1])-6), | |
| min(pix.width,int(b[2])+6),min(pix.height,int(b[3])+6))) | |
| crop.save(out/'assets'/f'{aid}.png') | |
| assets.append({'id':aid,'class':cls,'page':i+1,'bbox':b,'score':d['score'], | |
| 'w':crop.width,'h':crop.height,'aspect':round(crop.width/max(1,crop.height),2)}) | |
| print(f'LAYOUT_PAGE {i+1}/{n}', flush=True) | |
| (out/'assets.json').write_text(json.dumps({'dpi':200,'n_pages':n,'assets':assets,'counts':counters})) | |
| (out/'layout.json').write_text(json.dumps(layout)) | |
| print('LAYOUT_DONE',flush=True) | |
| if __name__ == '__main__': | |
| run(sys.argv[1],sys.argv[2]) | |