|
|
@@ -0,0 +1,126 @@
|
|
|
+
|
|
|
+import os
|
|
|
+import time
|
|
|
+
|
|
|
+import cv2
|
|
|
+import fitz
|
|
|
+import numpy as np
|
|
|
+import torch
|
|
|
+import unimernet.tasks as tasks
|
|
|
+import yaml
|
|
|
+from PIL import Image
|
|
|
+from torch.utils.data import DataLoader, Dataset
|
|
|
+from torchvision import transforms
|
|
|
+from ultralytics import YOLO
|
|
|
+from unimernet.common.config import Config
|
|
|
+from unimernet.processors import load_processor
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+class CustomPEKModel:
|
|
|
+ def __init__(self, ocr: bool = False, show_log: bool = False):
|
|
|
+ ## ======== model init ========##
|
|
|
+ with open('configs/model_configs.yaml') as f:
|
|
|
+ model_configs = yaml.load(f, Loader=yaml.FullLoader)
|
|
|
+ img_size = model_configs['model_args']['img_size']
|
|
|
+ conf_thres = model_configs['model_args']['conf_thres']
|
|
|
+ iou_thres = model_configs['model_args']['iou_thres']
|
|
|
+ device = model_configs['model_args']['device']
|
|
|
+ dpi = model_configs['model_args']['pdf_dpi']
|
|
|
+ mfd_model = mfd_model_init(model_configs['model_args']['mfd_weight'])
|
|
|
+ mfr_model, mfr_vis_processors = mfr_model_init(model_configs['model_args']['mfr_weight'], device=device)
|
|
|
+ mfr_transform = transforms.Compose([mfr_vis_processors, ])
|
|
|
+ layout_model = layout_model_init(model_configs['model_args']['layout_weight'])
|
|
|
+ ocr_model = ModifiedPaddleOCR(show_log=True)
|
|
|
+ print(now.strftime('%Y-%m-%d %H:%M:%S'))
|
|
|
+ print('Model init done!')
|
|
|
+ ## ======== model init ========##
|
|
|
+
|
|
|
+ def __call__(self, image):
|
|
|
+
|
|
|
+ # layout检测 + 公式检测
|
|
|
+ doc_layout_result = []
|
|
|
+ latex_filling_list = []
|
|
|
+ mf_image_list = []
|
|
|
+
|
|
|
+ img_H, img_W = image.shape[0], image.shape[1]
|
|
|
+ layout_res = layout_model(image, ignore_catids=[])
|
|
|
+ # 公式检测
|
|
|
+ mfd_res = mfd_model.predict(image, imgsz=img_size, conf=conf_thres, iou=iou_thres, verbose=True)[0]
|
|
|
+ for xyxy, conf, cla in zip(mfd_res.boxes.xyxy.cpu(), mfd_res.boxes.conf.cpu(), mfd_res.boxes.cls.cpu()):
|
|
|
+ xmin, ymin, xmax, ymax = [int(p.item()) for p in xyxy]
|
|
|
+ new_item = {
|
|
|
+ 'category_id': 13 + int(cla.item()),
|
|
|
+ 'poly': [xmin, ymin, xmax, ymin, xmax, ymax, xmin, ymax],
|
|
|
+ 'score': round(float(conf.item()), 2),
|
|
|
+ 'latex': '',
|
|
|
+ }
|
|
|
+ layout_res['layout_dets'].append(new_item)
|
|
|
+ latex_filling_list.append(new_item)
|
|
|
+ bbox_img = get_croped_image(Image.fromarray(image), [xmin, ymin, xmax, ymax])
|
|
|
+ mf_image_list.append(bbox_img)
|
|
|
+
|
|
|
+ layout_res['page_info'] = dict(
|
|
|
+ page_no=idx,
|
|
|
+ height=img_H,
|
|
|
+ width=img_W
|
|
|
+ )
|
|
|
+ doc_layout_result.append(layout_res)
|
|
|
+
|
|
|
+ # 公式识别,因为识别速度较慢,为了提速,把单个pdf的所有公式裁剪完,一起批量做识别。
|
|
|
+ a = time.time()
|
|
|
+ dataset = MathDataset(mf_image_list, transform=mfr_transform)
|
|
|
+ dataloader = DataLoader(dataset, batch_size=128, num_workers=0)
|
|
|
+ mfr_res = []
|
|
|
+ gpu_total_cost = 0
|
|
|
+ for imgs in dataloader:
|
|
|
+ imgs = imgs.to(device)
|
|
|
+ gpu_start = time.time()
|
|
|
+ output = mfr_model.generate({'image': imgs})
|
|
|
+ gpu_cost = time.time() - gpu_start
|
|
|
+ gpu_total_cost += gpu_cost
|
|
|
+ print(f"gpu_cost: {gpu_cost}")
|
|
|
+ mfr_res.extend(output['pred_str'])
|
|
|
+ print(f"gpu_total_cost: {gpu_total_cost}")
|
|
|
+ for res, latex in zip(latex_filling_list, mfr_res):
|
|
|
+ res['latex'] = latex_rm_whitespace(latex)
|
|
|
+ b = time.time()
|
|
|
+ print("formula nums:", len(mf_image_list), "mfr time:", round(b - a, 2))
|
|
|
+
|
|
|
+ # ocr识别
|
|
|
+ for idx, image in enumerate(img_list):
|
|
|
+ pil_img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
|
|
|
+ single_page_res = doc_layout_result[idx]['layout_dets']
|
|
|
+ single_page_mfdetrec_res = []
|
|
|
+ for res in single_page_res:
|
|
|
+ if int(res['category_id']) in [13, 14]:
|
|
|
+ xmin, ymin = int(res['poly'][0]), int(res['poly'][1])
|
|
|
+ xmax, ymax = int(res['poly'][4]), int(res['poly'][5])
|
|
|
+ single_page_mfdetrec_res.append({
|
|
|
+ "bbox": [xmin, ymin, xmax, ymax],
|
|
|
+ })
|
|
|
+ for res in single_page_res:
|
|
|
+ if int(res['category_id']) in [0, 1, 2, 4, 6, 7]: # 需要进行ocr的类别
|
|
|
+ xmin, ymin = int(res['poly'][0]), int(res['poly'][1])
|
|
|
+ xmax, ymax = int(res['poly'][4]), int(res['poly'][5])
|
|
|
+ crop_box = [xmin, ymin, xmax, ymax]
|
|
|
+ cropped_img = Image.new('RGB', pil_img.size, 'white')
|
|
|
+ cropped_img.paste(pil_img.crop(crop_box), crop_box)
|
|
|
+ cropped_img = cv2.cvtColor(np.asarray(cropped_img), cv2.COLOR_RGB2BGR)
|
|
|
+ ocr_res = ocr_model.ocr(cropped_img, mfd_res=single_page_mfdetrec_res)[0]
|
|
|
+ if ocr_res:
|
|
|
+ for box_ocr_res in ocr_res:
|
|
|
+ p1, p2, p3, p4 = box_ocr_res[0]
|
|
|
+ text, score = box_ocr_res[1]
|
|
|
+ doc_layout_result[idx]['layout_dets'].append({
|
|
|
+ 'category_id': 15,
|
|
|
+ 'poly': p1 + p2 + p3 + p4,
|
|
|
+ 'score': round(score, 2),
|
|
|
+ 'text': text,
|
|
|
+ })
|
|
|
+
|
|
|
+ output_dir = args.output
|
|
|
+ os.makedirs(output_dir, exist_ok=True)
|
|
|
+ basename = os.path.basename(single_pdf)[0:-4]
|
|
|
+ with open(os.path.join(output_dir, f'{basename}.json'), 'w') as f:
|
|
|
+ json.dump(doc_layout_result, f)
|