Liddlerain/ChatReviewer
0
1import fitz, io, os2from PIL import Image3from collections import Counter4import json5import re6 7class Paper:8 def __init__(self, path, title='', url='', abs='', authors=[]):9 # 初始化函数，根据pdf路径初始化Paper对象 10 self.url = url # 文章链接11 self.path = path # pdf路径12 self.section_names = [] # 段落标题13 self.section_texts = {} # 段落内容14 self.abs = abs15 self.title_page = 016 if title == '':17 self.pdf = fitz.open(self.path) # pdf文档18 self.title = self.get_title()19 self.parse_pdf()20 else:21 self.title = title22 self.authors = authors23 self.roman_num = ["I", "II", 'III', "IV", "V", "VI", "VII", "VIII", "IIX", "IX", "X"]24 self.digit_num = [str(d + 1) for d in range(10)]25 self.first_image = ''26 27 def parse_pdf(self):28 self.pdf = fitz.open(self.path) # pdf文档29 self.text_list = [page.get_text() for page in self.pdf]30 self.all_text = ' '.join(self.text_list)31 self.extract_section_infomation()32 self.section_texts.update({"title": self.title})33 self.pdf.close()34 35 # 定义一个函数，根据字体的大小，识别每个章节名称，并返回一个列表36 def get_chapter_names(self, ):37 # # 打开一个pdf文件38 doc = fitz.open(self.path) # pdf文档39 text_list = [page.get_text() for page in doc]40 all_text = ''41 for text in text_list:42 all_text += text43 # # 创建一个空列表，用于存储章节名称44 chapter_names = []45 for line in all_text.split('\n'):46 line_list = line.split(' ')47 if '.' in line:48 point_split_list = line.split('.')49 space_split_list = line.split(' ')50 if 1 < len(space_split_list) < 5:51 if 1 < len(point_split_list) < 5 and (52 point_split_list[0] in self.roman_num or point_split_list[0] in self.digit_num):53 # print("line:", line)54 chapter_names.append(line)55 56 return chapter_names57 58 def get_title(self):59 doc = self.pdf # 打开pdf文件60 max_font_size = 0 # 初始化最大字体大小为061 max_string = "" # 初始化最大字体大小对应的字符串为空62 max_font_sizes = [0]63 for page_index, page in enumerate(doc): # 遍历每一页64 text = page.get_text("dict") # 获取页面上的文本信息65 blocks = text["blocks"] # 获取文本块列表66 for block in blocks: # 遍历每个文本块67 if block["type"] == 0 and len(block['lines']): # 如果是文字类型68 if len(block["lines"][0]["spans"]):69 font_size = block["lines"][0]["spans"][0]["size"] # 获取第一行第一段文字的字体大小70 max_font_sizes.append(font_size)71 if font_size > max_font_size: # 如果字体大小大于当前最大值72 max_font_size = font_size # 更新最大值73 max_string = block["lines"][0]["spans"][0]["text"] # 更新最大值对应的字符串74 max_font_sizes.sort()75 # print("max_font_sizes", max_font_sizes[-10:])76 cur_title = ''77 for page_index, page in enumerate(doc): # 遍历每一页78 text = page.get_text("dict") # 获取页面上的文本信息79 blocks = text["blocks"] # 获取文本块列表80 for block in blocks: # 遍历每个文本块81 if block["type"] == 0 and len(block['lines']): # 如果是文字类型82 if len(block["lines"][0]["spans"]):83 cur_string = block["lines"][0]["spans"][0]["text"] # 更新最大值对应的字符串84 font_flags = block["lines"][0]["spans"][0]["flags"] # 获取第一行第一段文字的字体特征85 font_size = block["lines"][0]["spans"][0]["size"] # 获取第一行第一段文字的字体大小86 # print(font_size)87 if abs(font_size - max_font_sizes[-1]) < 0.3 or abs(font_size - max_font_sizes[-2]) < 0.3:88 # print("The string is bold.", max_string, "font_size:", font_size, "font_flags:", font_flags) 89 if len(cur_string) > 4 and "arXiv" not in cur_string:90 # print("The string is bold.", max_string, "font_size:", font_size, "font_flags:", font_flags) 91 if cur_title == '':92 cur_title += cur_string93 else:94 cur_title += ' ' + cur_string95 self.title_page = page_index96 # break97 title = cur_title.replace('\n', ' ')98 return title99 100 def extract_section_infomation(self):101 doc = fitz.open(self.path)102 103 # 获取文档中所有字体大小104 font_sizes = []105 for page in doc:106 blocks = page.get_text("dict")["blocks"]107 for block in blocks:108 if 'lines' not in block:109 continue110 lines = block["lines"]111 for line in lines:112 for span in line["spans"]:113 font_sizes.append(span["size"])114 most_common_size, _ = Counter(font_sizes).most_common(1)[0]115 116 # 按照最频繁的字体大小确定标题字体大小的阈值117 threshold = most_common_size * 1118 119 section_dict = {}120 last_heading = None121 subheadings = []122 heading_font = -1123 # 遍历每一页并查找子标题124 found_abstract = False125 upper_heading = False126 font_heading = False127 for page in doc:128 blocks = page.get_text("dict")["blocks"]129 for block in blocks:130 if not found_abstract:131 try:132 text = json.dumps(block)133 except:134 continue135 if re.search(r"\bAbstract\b", text, re.IGNORECASE):136 found_abstract = True137 last_heading = "Abstract"138 section_dict["Abstract"] = ""139 if found_abstract:140 if 'lines' not in block:141 continue142 lines = block["lines"]143 for line in lines:144 for span in line["spans"]:145 # 如果当前文本是子标题146 if not font_heading and span["text"].isupper() and sum(1 for c in span["text"] if c.isupper() and ('A' <= c <='Z')) > 4: # 针对一些标题大小一样,但是全大写的论文147 upper_heading = True148 heading = span["text"].strip()149 if "References" in heading: # reference 以后的内容不考虑150 self.section_names = subheadings151 self.section_texts = section_dict152 return153 subheadings.append(heading)154 if last_heading is not None:155 section_dict[last_heading] = section_dict[last_heading].strip()156 section_dict[heading] = ""157 last_heading = heading158 if not upper_heading and span["size"] > threshold and re.match( # 正常情况下,通过字体大小判断159 r"[A-Z][a-z]+(?:\s[A-Z][a-z]+)*",160 span["text"].strip()):161 font_heading = True162 if heading_font == -1:163 heading_font = span["size"]164 elif heading_font != span["size"]:165 continue166 heading = span["text"].strip()167 if "References" in heading: # reference 以后的内容不考虑168 self.section_names = subheadings169 self.section_texts = section_dict170 return171 subheadings.append(heading)172 if last_heading is not None:173 section_dict[last_heading] = section_dict[last_heading].strip()174 section_dict[heading] = ""175 last_heading = heading176 # 否则将当前文本添加到上一个子标题的文本中177 elif last_heading is not None:178 section_dict[last_heading] += " " + span["text"].strip()179 self.section_names = subheadings180 self.section_texts = section_dict181 182 183def main():184 path = r'demo.pdf'185 paper = Paper(path=path)186 paper.parse_pdf()187 # for key, value in paper.section_text_dict.items():188 # print(key, value)189 # print("*"*40)190 191 192if __name__ == '__main__':193 main()194