wehub-resource-sync
433 lines
16 KiB
Python
433 lines
16 KiB
Python
# coding=utf-8
|
|
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import base64
|
|
import mimetypes
|
|
import os
|
|
import random
|
|
import re
|
|
from io import BytesIO
|
|
|
|
import numpy as np
|
|
import requests
|
|
from packaging.version import Version
|
|
from PIL import Image, ImageDraw, ImageOps
|
|
|
|
from .image_utils import np2base64
|
|
from .log import logger
|
|
|
|
|
|
class DocParser(object):
|
|
"""DocParser"""
|
|
|
|
def __init__(self, ocr_lang="ch", layout_analysis=False, pdf_parser_config=None, use_gpu=None, device_id=None):
|
|
self.ocr_lang = ocr_lang
|
|
self.use_angle_cls = False
|
|
self.layout_analysis = layout_analysis
|
|
self.pdf_parser_config = pdf_parser_config
|
|
self.ocr_infer_model = None
|
|
self.use_gpu = use_gpu
|
|
self.device_id = device_id
|
|
|
|
def parse(self, doc, expand_to_a4_size=False, do_ocr=True):
|
|
"""
|
|
parse
|
|
"""
|
|
doc_type = mimetypes.guess_type(doc["doc"])[0]
|
|
|
|
if not doc_type or doc_type.startswith("image"):
|
|
image = self.read_image(doc["doc"])
|
|
elif doc_type == "application/pdf":
|
|
image = self.read_pdf(doc["doc"])
|
|
offset_x, offset_y = 0, 0
|
|
if expand_to_a4_size:
|
|
image, offset_x, offset_y = self.expand_image_to_a4_size(image, center=True)
|
|
img_h, img_w = image.shape[:2]
|
|
doc["image"] = np2base64(image)
|
|
doc["offset_x"] = offset_x
|
|
doc["offset_y"] = offset_y
|
|
doc["img_w"] = img_w
|
|
doc["img_h"] = img_h
|
|
if do_ocr:
|
|
ocr_result = self.ocr(image)
|
|
if expand_to_a4_size:
|
|
layout = []
|
|
for segment in ocr_result:
|
|
box = segment[0]
|
|
org_box = [
|
|
max(box[0] - offset_x, 0),
|
|
max(box[1] - offset_y, 0),
|
|
max(box[2] - offset_x, 0),
|
|
max(box[3] - offset_y, 0),
|
|
]
|
|
if len(segment) == 2:
|
|
layout.append((org_box, segment[1]))
|
|
elif len(segment) == 3:
|
|
layout.append((org_box, segment[1], segment[2]))
|
|
doc["layout"] = layout
|
|
else:
|
|
doc["layout"] = ocr_result
|
|
return doc
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
"""
|
|
Call parse
|
|
"""
|
|
return self.parse(*args, **kwargs)
|
|
|
|
def ocr(self, image, det=True, rec=True, cls=None):
|
|
"""
|
|
Call ocr for an image
|
|
"""
|
|
|
|
def _get_box(box):
|
|
box = [
|
|
min(box[0][0], box[3][0]), # x1
|
|
min(box[0][1], box[1][1]), # y1
|
|
max(box[1][0], box[2][0]), # x2
|
|
max(box[2][1], box[3][1]), # y2
|
|
]
|
|
return box
|
|
|
|
def _normal_box(box):
|
|
# Ensure the height and width of bbox are greater than zero
|
|
if box[3] - box[1] < 0 or box[2] - box[0] < 0:
|
|
return False
|
|
return True
|
|
|
|
def _is_ch(s):
|
|
for ch in s:
|
|
if "\u4e00" <= ch <= "\u9fff":
|
|
return True
|
|
return False
|
|
|
|
if self.ocr_infer_model is None:
|
|
self.init_ocr_inference()
|
|
if cls is None:
|
|
cls = self.use_angle_cls
|
|
remove = False if self.ppocr_version <= Version("2.6.0.1") else True
|
|
|
|
layout = []
|
|
if not self.layout_analysis:
|
|
ocr_result = self.ocr_infer_model.ocr(image, det, rec, cls)
|
|
ocr_result = ocr_result[0] if remove else ocr_result
|
|
for segment in ocr_result:
|
|
box = segment[0]
|
|
box = _get_box(box)
|
|
if not _normal_box(box):
|
|
continue
|
|
text = segment[1][0]
|
|
layout.append((box, text))
|
|
else:
|
|
layout_result = self.layout_analysis_engine(image)
|
|
for region in layout_result:
|
|
if region["type"] != "table":
|
|
ocr_result = region["res"]
|
|
for segment in ocr_result:
|
|
box = segment["text_region"]
|
|
box = _get_box(box)
|
|
if not _normal_box(box):
|
|
continue
|
|
text = segment["text"]
|
|
layout.append((box, text, region["type"]))
|
|
else:
|
|
bbox = region["bbox"]
|
|
table_result = region["res"]
|
|
html = table_result["html"]
|
|
cell_bbox = table_result["cell_bbox"]
|
|
table_list = []
|
|
lines = re.findall("<tr>(.*?)</tr>", html)
|
|
for line in lines:
|
|
table_list.extend(re.findall("<td.*?>(.*?)</td>", line))
|
|
for cell_box, text in zip(cell_bbox, table_list):
|
|
if self.ocr_lang == "ch":
|
|
box = [
|
|
bbox[0] + cell_box[0],
|
|
bbox[1] + cell_box[1],
|
|
bbox[0] + cell_box[4],
|
|
bbox[1] + cell_box[5],
|
|
]
|
|
else:
|
|
box = [
|
|
bbox[0] + cell_box[0],
|
|
bbox[1] + cell_box[1],
|
|
bbox[0] + cell_box[2],
|
|
bbox[1] + cell_box[3],
|
|
]
|
|
if not _normal_box(box):
|
|
continue
|
|
if _is_ch(text):
|
|
text = text.replace(" ", "")
|
|
layout.append((box, text, region["type"]))
|
|
return layout
|
|
|
|
@classmethod
|
|
def _get_buffer(self, data, file_like=False):
|
|
buff = None
|
|
if len(data) < 1024:
|
|
if os.path.exists(data):
|
|
buff = open(data, "rb").read()
|
|
elif data.startswith("http://") or data.startswith("https://"):
|
|
resp = requests.get(data, stream=True)
|
|
if not resp.ok:
|
|
raise RuntimeError("Failed to download the file from {}".format(data))
|
|
buff = resp.raw.read()
|
|
else:
|
|
raise FileNotFoundError("Image file {} not found!".format(data))
|
|
if buff is None:
|
|
buff = base64.b64decode(data)
|
|
if buff and file_like:
|
|
return BytesIO(buff)
|
|
return buff
|
|
|
|
@classmethod
|
|
def read_image(self, image):
|
|
"""
|
|
read image to np.ndarray
|
|
"""
|
|
image_buff = self._get_buffer(image)
|
|
|
|
# Use exif_transpose to correct orientation
|
|
_image = np.array(ImageOps.exif_transpose(Image.open(BytesIO(image_buff)).convert("RGB")))
|
|
return _image
|
|
|
|
@classmethod
|
|
def read_pdf(self, pdf, password=None):
|
|
"""
|
|
read pdf
|
|
"""
|
|
try:
|
|
import fitz
|
|
except ImportError:
|
|
raise RuntimeError(
|
|
"Need PyMuPDF to process pdf input. " "Please install module by: python3 -m pip install pymupdf"
|
|
)
|
|
if isinstance(pdf, fitz.Document):
|
|
return pdf
|
|
pdf_buff = self._get_buffer(pdf)
|
|
if not pdf_buff:
|
|
logger.warning("Failed to read pdf: %s...", pdf[:32])
|
|
return None
|
|
pdf_doc = fitz.Document(stream=pdf_buff, filetype="pdf")
|
|
if pdf_doc.needs_pass:
|
|
if pdf_doc.authenticate(password) == 0:
|
|
raise ValueError("The password of pdf is incorrect.")
|
|
|
|
if pdf_doc.page_count > 1:
|
|
logger.warning("Currently only parse the first page for PDF input with more than one page.")
|
|
|
|
page = pdf_doc.load_page(0)
|
|
# The original image is shrunk when convertd from PDF by fitz, so we scale the image size by 10 times
|
|
matrix = fitz.Matrix(10, 10)
|
|
image = np.array(self.get_page_image(page, matrix).convert("RGB"))
|
|
return image
|
|
|
|
@classmethod
|
|
def get_page_image(self, page, matrix):
|
|
"""
|
|
get page image
|
|
"""
|
|
pix = page.get_pixmap(matrix=matrix)
|
|
image_buff = pix.pil_tobytes("jpeg")
|
|
return Image.open(BytesIO(image_buff))
|
|
|
|
def init_ocr_inference(self):
|
|
"""
|
|
init ocr inference
|
|
"""
|
|
if self.ocr_infer_model is not None:
|
|
logger.warning("ocr model has already been initialized")
|
|
return
|
|
|
|
try:
|
|
import paddleocr
|
|
except ImportError:
|
|
raise RuntimeError(
|
|
"Need paddleocr to process image input. Please install module by: python3 -m pip install paddleocr"
|
|
)
|
|
self.ppocr_version = Version(paddleocr.__version__)
|
|
|
|
if not self.layout_analysis:
|
|
from paddleocr import PaddleOCR
|
|
|
|
self.ocr_infer_model = PaddleOCR(show_log=False, lang=self.ocr_lang)
|
|
else:
|
|
from paddleocr import PPStructure
|
|
|
|
self.layout_analysis_engine = PPStructure(table=True, ocr=True, show_log=False, lang=self.ocr_lang)
|
|
|
|
@classmethod
|
|
def _normalize_box(self, box, old_size, new_size, offset_x=0, offset_y=0):
|
|
"""normalize box"""
|
|
return [
|
|
int((box[0] + offset_x) * new_size[0] / old_size[0]),
|
|
int((box[1] + offset_y) * new_size[1] / old_size[1]),
|
|
int((box[2] + offset_x) * new_size[0] / old_size[0]),
|
|
int((box[3] + offset_y) * new_size[1] / old_size[1]),
|
|
]
|
|
|
|
@classmethod
|
|
def expand_image_to_a4_size(self, image, center=False):
|
|
"""expand image to a4 size"""
|
|
h, w = image.shape[:2]
|
|
offset_x, offset_y = 0, 0
|
|
if h * 1.0 / w >= 1.42:
|
|
exp_w = int(h / 1.414 - w)
|
|
if center:
|
|
offset_x = int(exp_w / 2)
|
|
exp_img = np.zeros((h, offset_x, 3), dtype="uint8")
|
|
exp_img.fill(255)
|
|
image = np.hstack([exp_img, image, exp_img])
|
|
else:
|
|
exp_img = np.zeros((h, exp_w, 3), dtype="uint8")
|
|
exp_img.fill(255)
|
|
image = np.hstack([image, exp_img])
|
|
elif h * 1.0 / w <= 1.40:
|
|
exp_h = int(w * 1.414 - h)
|
|
if center:
|
|
offset_y = int(exp_h / 2)
|
|
exp_img = np.zeros((offset_y, w, 3), dtype="uint8")
|
|
exp_img.fill(255)
|
|
image = np.vstack([exp_img, image, exp_img])
|
|
else:
|
|
exp_img = np.zeros((exp_h, w, 3), dtype="uint8")
|
|
exp_img.fill(255)
|
|
image = np.vstack([image, exp_img])
|
|
return image, offset_x, offset_y
|
|
|
|
@classmethod
|
|
def write_image_with_results(
|
|
self, image, layout=None, result=None, save_path=None, return_image=False, format=None, max_size=None
|
|
):
|
|
"""
|
|
write image with boxes and results
|
|
"""
|
|
|
|
def _flatten_results(results):
|
|
"""flatten results"""
|
|
is_single = False
|
|
if not isinstance(results, list):
|
|
results = [results]
|
|
is_single = True
|
|
flat_results = []
|
|
|
|
def _flatten(result):
|
|
flat_result = []
|
|
for key, vals in result.items():
|
|
for val in vals:
|
|
new_val = val.copy()
|
|
if val.get("relations"):
|
|
new_val["relations"] = _flatten(val["relations"])
|
|
new_val["label"] = key
|
|
flat_result.append(new_val)
|
|
return flat_result
|
|
|
|
for result in results:
|
|
flat_results.append(_flatten(result))
|
|
if is_single:
|
|
return flat_results[0]
|
|
return flat_results
|
|
|
|
def _write_results(results, color=None, root=True, parent_centers=None):
|
|
for segment in results:
|
|
if "bbox" not in segment.keys():
|
|
continue
|
|
boxes = segment["bbox"]
|
|
if not isinstance(boxes[0], list):
|
|
boxes = [boxes]
|
|
centers = []
|
|
plot_boxes = []
|
|
for box in boxes:
|
|
x1, y1, x2, y2 = box
|
|
plot_box = [
|
|
(x1, y1),
|
|
(x2, y1),
|
|
(x2, y2),
|
|
(x1, y2),
|
|
]
|
|
plot_boxes.append(plot_box)
|
|
centers.append(((x2 - x1) / 2 + x1, (y2 - y1) / 2 + y1))
|
|
if root:
|
|
while True:
|
|
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
|
|
if sum(color) < 480:
|
|
break
|
|
for box in plot_boxes:
|
|
draw_render.polygon(box, fill=color)
|
|
if parent_centers:
|
|
for p_c in parent_centers:
|
|
for c in centers:
|
|
draw_render.line((p_c[0], p_c[1], c[0], c[1]), fill=125, width=3)
|
|
if isinstance(segment, dict) and segment.get("relations"):
|
|
_write_results(segment["relations"], color, root=False, parent_centers=centers)
|
|
|
|
random.seed(0)
|
|
_image = self.read_image(image)
|
|
_image = Image.fromarray(np.uint8(_image))
|
|
h, w = _image.height, _image.width
|
|
img_render = _image.copy()
|
|
draw_render = ImageDraw.Draw(img_render)
|
|
|
|
if layout:
|
|
for segment in layout:
|
|
if isinstance(segment, dict):
|
|
box = segment["bbox"]
|
|
else:
|
|
box = segment[0]
|
|
box = [
|
|
(box[0], box[1]),
|
|
(box[2], box[1]),
|
|
(box[2], box[3]),
|
|
(box[0], box[3]),
|
|
]
|
|
while True:
|
|
color = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
|
|
if sum(color) < 480:
|
|
break
|
|
draw_render.polygon(box, fill=color)
|
|
|
|
elif result:
|
|
flatten_results = _flatten_results(result)
|
|
_write_results(flatten_results, color=None, root=True)
|
|
|
|
img_render = Image.blend(_image, img_render, 0.3)
|
|
img_show = Image.new("RGB", (w, h), (255, 255, 255))
|
|
img_show.paste(img_render, (0, 0, w, h))
|
|
w, h = img_show.width, img_show.height
|
|
if max_size and max(w, h) > max_size:
|
|
if max(w, h) == h:
|
|
new_size = (int(w * max_size / h), max_size)
|
|
else:
|
|
new_size = (max_size, int(h * max_size / w))
|
|
img_show = img_show.resize(new_size)
|
|
|
|
if save_path:
|
|
dir_path = os.path.dirname(save_path)
|
|
if dir_path and not os.path.isdir(dir_path):
|
|
os.makedirs(dir_path)
|
|
img_show.save(save_path)
|
|
if return_image:
|
|
return np.array(img_show)
|
|
elif return_image:
|
|
return np.array(img_show)
|
|
else:
|
|
buff = BytesIO()
|
|
if format is None:
|
|
format = "jpeg"
|
|
if format.lower() == "jpg":
|
|
format = "jpeg"
|
|
img_show.save(buff, format=format, quality=90)
|
|
return buff
|