import io import cv2 from PIL import Image import numpy as np from format_convert.convert_image import image_preprocess from format_convert.utils import add_div, judge_error_code, get_table_html, sort_object class _Document: def __init__(self, doc_path): self.doc_path = doc_path # Document's child -> Page self.children = [] self.error_code = None def add_child(self, child): if child.error_code is None: self.children.append(child) else: self.error_code = child.error_code def get_html(self): if self.error_code is not None: return self.error_code html_text = "" for child in self.children: # 先调用get_html才能更新error_code child_html_text = child.get_html() print("Document", self.error_code, child.error_code, type(child), child.page_no) if child.error_code is not None: self.error_code = child.error_code return self.error_code else: html_text += child_html_text return [html_text] class _Page: def __init__(self, page, page_no): self.page = page self.page_no = page_no # Page's child -> Image, Table, Sentence self.children = [] self.error_code = None # objs in tables self.in_table_objs = set() def add_child(self, child): if child.error_code is None: self.children.append(child) else: self.error_code = child.error_code def get_html(self): if self.error_code is not None: return "" html_text = "" self.children = sort_object(self.children) for child in self.children: # 先调用get_html才能更新error_code child_html_text = child.get_html() print("Page", self.error_code, child.error_code, type(child)) if child.error_code is not None: self.error_code = child.error_code return "" else: html_text += child_html_text return html_text class _Image: def __init__(self, content, path): self.content = content self.path = path # 来源 self.is_from_pdf = False # 位置 self.x = 0 self.y = 0 # 识别结果 self.otr_result = None self.ocr_result = None # Image's child -> Table, Sentence self.children = [] self.error_code = None # objs in tables self.in_table_objs = set() def add_child(self, child): if child.error_code is None: self.children.append(child) else: self.error_code = child.error_code def get_html(self): # 将Image转为Sentence,table self.convert() if self.error_code is not None: return "" html_text = "" self.children = sort_object(self.children) for child in self.children: # 先调用get_html才能更新error_code child_html_text = child.get_html() print("Image", self.error_code, child.error_code, type(child)) if child.error_code is not None: self.error_code = child.error_code return "" else: html_text += child_html_text return html_text def get_text(self): return def convert(self): # 二进制转numpy image_np = Image.open(io.BytesIO(self.content)) image_np = cv2.cvtColor(np.asarray(image_np), cv2.COLOR_RGB2BGR) obj_list = image_preprocess(image_np, self.path, use_ocr=True) if judge_error_code(obj_list): self.error_code = obj_list return for obj in obj_list: self.add_child(obj) class _Table: def __init__(self, content, bbox): self.content = content self.bbox = bbox self.x = bbox[0] self.y = bbox[1] self.shape = (len(content), len(content[0])) self.error_code = None def get_html(self): if self.error_code is not None: return "" # 将二维数组转为html table html_text = get_table_html(self.content) return html_text class _Sentence: def __init__(self, content, bbox): self.content = content # 位置 self.bbox = bbox self.x = bbox[0] self.y = bbox[1] self.error_code = None def get_html(self): if self.error_code is not None: return "" print("===========_Sentence get_html", self.content) return add_div(self.content) class TextBox: def __init__(self, bbox, text): self.bbox = bbox self.text = text def get_text(self): return self.text class TableLine: def __init__(self, bbox): self.bbox = bbox