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)