Spaces:
Running on Zero
Running on Zero
File size: 4,234 Bytes
c58c39d | 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 | """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])
|