pdf2text_recogPara_v2.py 127 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545154615471548154915501551155215531554155515561557155815591560156115621563156415651566156715681569157015711572157315741575157615771578157915801581158215831584158515861587158815891590159115921593159415951596159715981599160016011602160316041605160616071608160916101611161216131614161516161617161816191620162116221623162416251626162716281629163016311632163316341635163616371638163916401641164216431644164516461647164816491650165116521653165416551656165716581659166016611662166316641665166616671668166916701671167216731674167516761677167816791680168116821683168416851686168716881689169016911692169316941695169616971698169917001701170217031704170517061707170817091710171117121713171417151716171717181719172017211722172317241725172617271728172917301731173217331734173517361737173817391740174117421743174417451746174717481749175017511752175317541755175617571758175917601761176217631764176517661767176817691770177117721773177417751776177717781779178017811782178317841785178617871788178917901791179217931794179517961797179817991800180118021803180418051806180718081809181018111812181318141815181618171818181918201821182218231824182518261827182818291830183118321833183418351836183718381839184018411842184318441845184618471848184918501851185218531854185518561857185818591860186118621863186418651866186718681869187018711872187318741875187618771878187918801881188218831884188518861887188818891890189118921893189418951896189718981899190019011902190319041905190619071908190919101911191219131914191519161917191819191920192119221923192419251926192719281929193019311932193319341935193619371938193919401941194219431944194519461947194819491950195119521953195419551956195719581959196019611962196319641965196619671968196919701971197219731974197519761977197819791980198119821983198419851986198719881989199019911992199319941995199619971998199920002001200220032004200520062007200820092010201120122013201420152016201720182019202020212022202320242025202620272028202920302031203220332034203520362037203820392040204120422043204420452046204720482049205020512052205320542055205620572058205920602061206220632064206520662067206820692070207120722073207420752076207720782079208020812082208320842085208620872088208920902091209220932094209520962097209820992100210121022103210421052106210721082109211021112112211321142115211621172118211921202121212221232124212521262127212821292130213121322133213421352136213721382139214021412142214321442145214621472148214921502151215221532154215521562157215821592160216121622163216421652166216721682169217021712172217321742175217621772178217921802181218221832184218521862187218821892190219121922193219421952196219721982199220022012202220322042205220622072208220922102211221222132214221522162217221822192220222122222223222422252226222722282229223022312232223322342235223622372238223922402241224222432244224522462247224822492250225122522253225422552256225722582259226022612262226322642265226622672268226922702271227222732274227522762277227822792280228122822283228422852286228722882289229022912292229322942295229622972298229923002301230223032304230523062307230823092310231123122313231423152316231723182319232023212322232323242325232623272328232923302331233223332334233523362337233823392340234123422343234423452346234723482349235023512352235323542355235623572358235923602361236223632364236523662367236823692370237123722373237423752376237723782379238023812382238323842385238623872388238923902391239223932394239523962397239823992400240124022403240424052406240724082409241024112412241324142415241624172418241924202421242224232424242524262427242824292430243124322433243424352436243724382439244024412442244324442445244624472448244924502451245224532454245524562457245824592460246124622463246424652466246724682469247024712472247324742475247624772478247924802481248224832484248524862487248824892490249124922493249424952496249724982499250025012502250325042505250625072508250925102511251225132514251525162517251825192520252125222523252425252526252725282529253025312532253325342535253625372538253925402541254225432544254525462547254825492550255125522553255425552556255725582559256025612562256325642565256625672568256925702571257225732574257525762577257825792580258125822583258425852586258725882589259025912592259325942595259625972598259926002601260226032604260526062607260826092610261126122613261426152616261726182619262026212622262326242625262626272628262926302631263226332634263526362637263826392640264126422643264426452646264726482649265026512652265326542655265626572658265926602661266226632664266526662667266826692670267126722673267426752676267726782679268026812682268326842685268626872688268926902691269226932694269526962697269826992700270127022703270427052706270727082709271027112712271327142715271627172718271927202721272227232724272527262727272827292730273127322733273427352736273727382739274027412742274327442745274627472748274927502751275227532754275527562757275827592760276127622763276427652766276727682769277027712772277327742775277627772778277927802781278227832784278527862787278827892790279127922793279427952796279727982799280028012802280328042805280628072808280928102811281228132814281528162817281828192820282128222823282428252826282728282829283028312832283328342835283628372838283928402841284228432844284528462847284828492850285128522853285428552856285728582859286028612862286328642865286628672868286928702871287228732874287528762877287828792880288128822883288428852886288728882889289028912892289328942895289628972898289929002901290229032904290529062907290829092910291129122913291429152916291729182919292029212922292329242925292629272928292929302931293229332934293529362937293829392940294129422943294429452946294729482949295029512952295329542955295629572958295929602961296229632964296529662967296829692970297129722973297429752976297729782979298029812982298329842985298629872988298929902991299229932994299529962997299829993000300130023003300430053006300730083009301030113012301330143015301630173018301930203021302230233024302530263027302830293030303130323033303430353036303730383039304030413042304330443045304630473048304930503051305230533054305530563057305830593060306130623063306430653066306730683069307030713072307330743075307630773078307930803081308230833084308530863087308830893090309130923093309430953096309730983099310031013102310331043105310631073108310931103111311231133114311531163117311831193120312131223123312431253126312731283129313031313132313331343135313631373138313931403141314231433144314531463147314831493150315131523153315431553156315731583159316031613162316331643165316631673168316931703171317231733174317531763177317831793180318131823183318431853186318731883189319031913192319331943195319631973198319932003201320232033204320532063207320832093210321132123213321432153216321732183219322032213222322332243225322632273228322932303231323232333234323532363237323832393240324132423243324432453246324732483249325032513252325332543255325632573258325932603261326232633264326532663267326832693270327132723273327432753276327732783279328032813282328332843285328632873288328932903291329232933294329532963297329832993300330133023303330433053306330733083309331033113312331333143315331633173318331933203321332233233324332533263327332833293330333133323333333433353336333733383339334033413342334333443345334633473348334933503351335233533354335533563357335833593360336133623363336433653366336733683369337033713372337333743375337633773378337933803381338233833384338533863387338833893390339133923393339433953396339733983399340034013402340334043405340634073408340934103411341234133414341534163417341834193420342134223423342434253426342734283429343034313432343334343435343634373438343934403441344234433444344534463447344834493450345134523453345434553456345734583459346034613462346334643465346634673468346934703471347234733474347534763477347834793480348134823483348434853486348734883489349034913492349334943495349634973498349935003501350235033504350535063507350835093510351135123513351435153516351735183519352035213522352335243525352635273528352935303531353235333534353535363537353835393540354135423543354435453546
  1. import os
  2. import sys
  3. import json
  4. import re
  5. import math
  6. import unicodedata
  7. from collections import Counter
  8. import numpy as np
  9. from termcolor import cprint
  10. from magic_pdf.libs import fitz
  11. from magic_pdf.libs import NLPModels
  12. if sys.version_info[0] >= 3:
  13. sys.stdout.reconfigure(encoding="utf-8") # type: ignore
  14. def open_pdf(pdf_path):
  15. try:
  16. pdf_document = fitz.open(pdf_path) # type: ignore
  17. return pdf_document
  18. except Exception as e:
  19. print(f"无法打开PDF文件:{pdf_path}。原因是:{e}")
  20. raise e
  21. def print_green_on_red(text):
  22. cprint(text, "green", "on_red", attrs=["bold"], end="\n\n")
  23. def print_green(text):
  24. print()
  25. cprint(text, "green", attrs=["bold"], end="\n\n")
  26. def print_red(text):
  27. print()
  28. cprint(text, "red", attrs=["bold"], end="\n\n")
  29. def print_yellow(text):
  30. print()
  31. cprint(text, "yellow", attrs=["bold"], end="\n\n")
  32. def safe_get(dict_obj, key, default):
  33. val = dict_obj.get(key)
  34. if val is None:
  35. return default
  36. else:
  37. return val
  38. def is_bbox_overlap(bbox1, bbox2):
  39. """
  40. This function checks if bbox1 and bbox2 overlap or not
  41. Parameters
  42. ----------
  43. bbox1 : list
  44. bbox1
  45. bbox2 : list
  46. bbox2
  47. Returns
  48. -------
  49. bool
  50. True if bbox1 and bbox2 overlap, else False
  51. """
  52. x0_1, y0_1, x1_1, y1_1 = bbox1
  53. x0_2, y0_2, x1_2, y1_2 = bbox2
  54. if x0_1 > x1_2 or x0_2 > x1_1:
  55. return False
  56. if y0_1 > y1_2 or y0_2 > y1_1:
  57. return False
  58. return True
  59. def is_in_bbox(bbox1, bbox2):
  60. """
  61. This function checks if bbox1 is in bbox2
  62. Parameters
  63. ----------
  64. bbox1 : list
  65. bbox1
  66. bbox2 : list
  67. bbox2
  68. Returns
  69. -------
  70. bool
  71. True if bbox1 is in bbox2, else False
  72. """
  73. x0_1, y0_1, x1_1, y1_1 = bbox1
  74. x0_2, y0_2, x1_2, y1_2 = bbox2
  75. if x0_1 >= x0_2 and y0_1 >= y0_2 and x1_1 <= x1_2 and y1_1 <= y1_2:
  76. return True
  77. else:
  78. return False
  79. def calculate_para_bbox(lines):
  80. """
  81. This function calculates the minimum bbox of the paragraph
  82. Parameters
  83. ----------
  84. lines : list
  85. lines
  86. Returns
  87. -------
  88. para_bbox : list
  89. bbox of the paragraph
  90. """
  91. x0 = min(line["bbox"][0] for line in lines)
  92. y0 = min(line["bbox"][1] for line in lines)
  93. x1 = max(line["bbox"][2] for line in lines)
  94. y1 = max(line["bbox"][3] for line in lines)
  95. return [x0, y0, x1, y1]
  96. def is_line_right_aligned_from_neighbors(curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, direction=2):
  97. """
  98. This function checks if the line is right aligned from its neighbors
  99. Parameters
  100. ----------
  101. curr_line_bbox : list
  102. bbox of the current line
  103. prev_line_bbox : list
  104. bbox of the previous line
  105. next_line_bbox : list
  106. bbox of the next line
  107. avg_char_width : float
  108. average of char widths
  109. direction : int
  110. 0 for prev, 1 for next, 2 for both
  111. Returns
  112. -------
  113. bool
  114. True if the line is right aligned from its neighbors, False otherwise.
  115. """
  116. horizontal_ratio = 0.5
  117. horizontal_thres = horizontal_ratio * avg_char_width
  118. _, _, x1, _ = curr_line_bbox
  119. _, _, prev_x1, _ = prev_line_bbox if prev_line_bbox else (0, 0, 0, 0)
  120. _, _, next_x1, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
  121. if direction == 0:
  122. return abs(x1 - prev_x1) < horizontal_thres
  123. elif direction == 1:
  124. return abs(x1 - next_x1) < horizontal_thres
  125. elif direction == 2:
  126. return abs(x1 - prev_x1) < horizontal_thres and abs(x1 - next_x1) < horizontal_thres
  127. else:
  128. return False
  129. def is_line_left_aligned_from_neighbors(curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, direction=2):
  130. """
  131. This function checks if the line is left aligned from its neighbors
  132. Parameters
  133. ----------
  134. curr_line_bbox : list
  135. bbox of the current line
  136. prev_line_bbox : list
  137. bbox of the previous line
  138. next_line_bbox : list
  139. bbox of the next line
  140. avg_char_width : float
  141. average of char widths
  142. direction : int
  143. 0 for prev, 1 for next, 2 for both
  144. Returns
  145. -------
  146. bool
  147. True if the line is left aligned from its neighbors, False otherwise.
  148. """
  149. horizontal_ratio = 0.5
  150. horizontal_thres = horizontal_ratio * avg_char_width
  151. x0, _, _, _ = curr_line_bbox
  152. prev_x0, _, _, _ = prev_line_bbox if prev_line_bbox else (0, 0, 0, 0)
  153. next_x0, _, _, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
  154. if direction == 0:
  155. return abs(x0 - prev_x0) < horizontal_thres
  156. elif direction == 1:
  157. return abs(x0 - next_x0) < horizontal_thres
  158. elif direction == 2:
  159. return abs(x0 - prev_x0) < horizontal_thres and abs(x0 - next_x0) < horizontal_thres
  160. else:
  161. return False
  162. def end_with_punctuation(line_text):
  163. """
  164. This function checks if the line ends with punctuation marks
  165. """
  166. english_end_puncs = [".", "?", "!"]
  167. chinese_end_puncs = ["。", "?", "!"]
  168. end_puncs = english_end_puncs + chinese_end_puncs
  169. last_non_space_char = None
  170. for ch in line_text[::-1]:
  171. if not ch.isspace():
  172. last_non_space_char = ch
  173. break
  174. if last_non_space_char is None:
  175. return False
  176. return last_non_space_char in end_puncs
  177. def is_nested_list(lst):
  178. if isinstance(lst, list):
  179. return any(isinstance(sub, list) for sub in lst)
  180. return False
  181. class DenseSingleLineBlockException(Exception):
  182. """
  183. This class defines the exception type for dense single line-block.
  184. """
  185. def __init__(self, message="DenseSingleLineBlockException"):
  186. self.message = message
  187. super().__init__(self.message)
  188. def __str__(self):
  189. return f"{self.message}"
  190. def __repr__(self):
  191. return f"{self.message}"
  192. class TitleDetectionException(Exception):
  193. """
  194. This class defines the exception type for title detection.
  195. """
  196. def __init__(self, message="TitleDetectionException"):
  197. self.message = message
  198. super().__init__(self.message)
  199. def __str__(self):
  200. return f"{self.message}"
  201. def __repr__(self):
  202. return f"{self.message}"
  203. class TitleLevelException(Exception):
  204. """
  205. This class defines the exception type for title level.
  206. """
  207. def __init__(self, message="TitleLevelException"):
  208. self.message = message
  209. super().__init__(self.message)
  210. def __str__(self):
  211. return f"{self.message}"
  212. def __repr__(self):
  213. return f"{self.message}"
  214. class ParaSplitException(Exception):
  215. """
  216. This class defines the exception type for paragraph splitting.
  217. """
  218. def __init__(self, message="ParaSplitException"):
  219. self.message = message
  220. super().__init__(self.message)
  221. def __str__(self):
  222. return f"{self.message}"
  223. def __repr__(self):
  224. return f"{self.message}"
  225. class ParaMergeException(Exception):
  226. """
  227. This class defines the exception type for paragraph merging.
  228. """
  229. def __init__(self, message="ParaMergeException"):
  230. self.message = message
  231. super().__init__(self.message)
  232. def __str__(self):
  233. return f"{self.message}"
  234. def __repr__(self):
  235. return f"{self.message}"
  236. class DiscardByException:
  237. """
  238. This class discards pdf files by exception
  239. """
  240. def __init__(self) -> None:
  241. pass
  242. def discard_by_single_line_block(self, pdf_dic, exception: DenseSingleLineBlockException):
  243. """
  244. This function discards pdf files by single line block exception
  245. Parameters
  246. ----------
  247. pdf_dic : dict
  248. pdf dictionary
  249. exception : str
  250. exception message
  251. Returns
  252. -------
  253. error_message : str
  254. """
  255. exception_page_nums = 0
  256. page_num = 0
  257. for page_id, page in pdf_dic.items():
  258. if page_id.startswith("page_"):
  259. page_num += 1
  260. if "preproc_blocks" in page.keys():
  261. preproc_blocks = page["preproc_blocks"]
  262. all_single_line_blocks = []
  263. for block in preproc_blocks:
  264. if len(block["lines"]) == 1:
  265. all_single_line_blocks.append(block)
  266. if len(preproc_blocks) > 0 and len(all_single_line_blocks) / len(preproc_blocks) > 0.9:
  267. exception_page_nums += 1
  268. if page_num == 0:
  269. return None
  270. if exception_page_nums / page_num > 0.1: # Low ratio means basically, whenever this is the case, it is discarded
  271. return exception.message
  272. return None
  273. def discard_by_title_detection(self, pdf_dic, exception: TitleDetectionException):
  274. """
  275. This function discards pdf files by title detection exception
  276. Parameters
  277. ----------
  278. pdf_dic : dict
  279. pdf dictionary
  280. exception : str
  281. exception message
  282. Returns
  283. -------
  284. error_message : str
  285. """
  286. # return exception.message
  287. return None
  288. def discard_by_title_level(self, pdf_dic, exception: TitleLevelException):
  289. """
  290. This function discards pdf files by title level exception
  291. Parameters
  292. ----------
  293. pdf_dic : dict
  294. pdf dictionary
  295. exception : str
  296. exception message
  297. Returns
  298. -------
  299. error_message : str
  300. """
  301. # return exception.message
  302. return None
  303. def discard_by_split_para(self, pdf_dic, exception: ParaSplitException):
  304. """
  305. This function discards pdf files by split para exception
  306. Parameters
  307. ----------
  308. pdf_dic : dict
  309. pdf dictionary
  310. exception : str
  311. exception message
  312. Returns
  313. -------
  314. error_message : str
  315. """
  316. # return exception.message
  317. return None
  318. def discard_by_merge_para(self, pdf_dic, exception: ParaMergeException):
  319. """
  320. This function discards pdf files by merge para exception
  321. Parameters
  322. ----------
  323. pdf_dic : dict
  324. pdf dictionary
  325. exception : str
  326. exception message
  327. Returns
  328. -------
  329. error_message : str
  330. """
  331. # return exception.message
  332. return None
  333. class LayoutFilterProcessor:
  334. def __init__(self) -> None:
  335. pass
  336. def batch_process_blocks(self, pdf_dict):
  337. """
  338. This function processes the blocks in batch.
  339. Parameters
  340. ----------
  341. self : object
  342. The instance of the class.
  343. pdf_dict : dict
  344. pdf dictionary
  345. Returns
  346. -------
  347. pdf_dict : dict
  348. pdf dictionary
  349. """
  350. for page_id, blocks in pdf_dict.items():
  351. if page_id.startswith("page_"):
  352. if "layout_bboxes" in blocks.keys() and "para_blocks" in blocks.keys():
  353. layout_bbox_objs = blocks["layout_bboxes"]
  354. if layout_bbox_objs is None:
  355. continue
  356. layout_bboxes = [bbox_obj["layout_bbox"] for bbox_obj in layout_bbox_objs]
  357. # Enlarge each value of x0, y0, x1, y1 for each layout_bbox to prevent loss of text.
  358. layout_bboxes = [
  359. [math.ceil(x0), math.ceil(y0), math.ceil(x1), math.ceil(y1)] for x0, y0, x1, y1 in layout_bboxes
  360. ]
  361. para_blocks = blocks["para_blocks"]
  362. if para_blocks is None:
  363. continue
  364. for lb_bbox in layout_bboxes:
  365. for i, para_block in enumerate(para_blocks):
  366. para_bbox = para_block["bbox"]
  367. para_blocks[i]["in_layout"] = 0
  368. if is_in_bbox(para_bbox, lb_bbox):
  369. para_blocks[i]["in_layout"] = 1
  370. blocks["para_blocks"] = para_blocks
  371. return pdf_dict
  372. class RawBlockProcessor:
  373. def __init__(self) -> None:
  374. self.y_tolerance = 2
  375. self.pdf_dic = {}
  376. def __span_flags_decomposer(self, span_flags):
  377. """
  378. Make font flags human readable.
  379. Parameters
  380. ----------
  381. self : object
  382. The instance of the class.
  383. span_flags : int
  384. span flags
  385. Returns
  386. -------
  387. l : dict
  388. decomposed flags
  389. """
  390. l = {
  391. "is_superscript": False,
  392. "is_italic": False,
  393. "is_serifed": False,
  394. "is_sans_serifed": False,
  395. "is_monospaced": False,
  396. "is_proportional": False,
  397. "is_bold": False,
  398. }
  399. if span_flags & 2**0:
  400. l["is_superscript"] = True # 表示上标
  401. if span_flags & 2**1:
  402. l["is_italic"] = True # 表示斜体
  403. if span_flags & 2**2:
  404. l["is_serifed"] = True # 表示衬线字体
  405. else:
  406. l["is_sans_serifed"] = True # 表示非衬线字体
  407. if span_flags & 2**3:
  408. l["is_monospaced"] = True # 表示等宽字体
  409. else:
  410. l["is_proportional"] = True # 表示比例字体
  411. if span_flags & 2**4:
  412. l["is_bold"] = True # 表示粗体
  413. return l
  414. def __make_new_lines(self, raw_lines):
  415. """
  416. This function makes new lines.
  417. Parameters
  418. ----------
  419. self : object
  420. The instance of the class.
  421. raw_lines : list
  422. raw lines
  423. Returns
  424. -------
  425. new_lines : list
  426. new lines
  427. """
  428. new_lines = []
  429. new_line = None
  430. for raw_line in raw_lines:
  431. raw_line_bbox = raw_line["bbox"]
  432. raw_line_spans = raw_line["spans"]
  433. raw_line_text = "".join([span["text"] for span in raw_line_spans])
  434. raw_line_dir = raw_line.get("dir", None)
  435. decomposed_line_spans = []
  436. for span in raw_line_spans:
  437. raw_flags = span["flags"]
  438. decomposed_flags = self.__span_flags_decomposer(raw_flags)
  439. span["decomposed_flags"] = decomposed_flags
  440. decomposed_line_spans.append(span)
  441. if new_line is None: # Handle the first line
  442. new_line = {
  443. "bbox": raw_line_bbox,
  444. "text": raw_line_text,
  445. "dir": raw_line_dir if raw_line_dir else (0, 0),
  446. "spans": decomposed_line_spans,
  447. }
  448. else: # Handle the rest lines
  449. if (
  450. abs(raw_line_bbox[1] - new_line["bbox"][1]) <= self.y_tolerance
  451. and abs(raw_line_bbox[3] - new_line["bbox"][3]) <= self.y_tolerance
  452. ):
  453. new_line["bbox"] = (
  454. min(new_line["bbox"][0], raw_line_bbox[0]), # left
  455. new_line["bbox"][1], # top
  456. max(new_line["bbox"][2], raw_line_bbox[2]), # right
  457. raw_line_bbox[3], # bottom
  458. )
  459. new_line["text"] += raw_line_text
  460. new_line["spans"].extend(raw_line_spans)
  461. new_line["dir"] = (
  462. new_line["dir"][0] + raw_line_dir[0],
  463. new_line["dir"][1] + raw_line_dir[1],
  464. )
  465. else:
  466. new_lines.append(new_line)
  467. new_line = {
  468. "bbox": raw_line_bbox,
  469. "text": raw_line_text,
  470. "dir": raw_line_dir if raw_line_dir else (0, 0),
  471. "spans": raw_line_spans,
  472. }
  473. if new_line:
  474. new_lines.append(new_line)
  475. return new_lines
  476. def __make_new_block(self, raw_block):
  477. """
  478. This function makes a new block.
  479. Parameters
  480. ----------
  481. self : object
  482. The instance of the class.
  483. ----------
  484. raw_block : dict
  485. a raw block
  486. Returns
  487. -------
  488. new_block : dict
  489. """
  490. new_block = {}
  491. block_id = raw_block["number"]
  492. block_bbox = raw_block["bbox"]
  493. block_text = "".join(span["text"] for line in raw_block["lines"] for span in line["spans"])
  494. raw_lines = raw_block["lines"]
  495. block_lines = self.__make_new_lines(raw_lines)
  496. new_block["block_id"] = block_id
  497. new_block["bbox"] = block_bbox
  498. new_block["text"] = block_text
  499. new_block["lines"] = block_lines
  500. return new_block
  501. def batch_process_blocks(self, pdf_dic):
  502. """
  503. This function processes the blocks in batch.
  504. Parameters
  505. ----------
  506. self : object
  507. The instance of the class.
  508. ----------
  509. blocks : list
  510. Input block is a list of raw blocks.
  511. Returns
  512. -------
  513. result_dict : dict
  514. result dictionary
  515. """
  516. for page_id, blocks in pdf_dic.items():
  517. if page_id.startswith("page_"):
  518. para_blocks = []
  519. if "preproc_blocks" in blocks.keys():
  520. input_blocks = blocks["preproc_blocks"]
  521. for raw_block in input_blocks:
  522. new_block = self.__make_new_block(raw_block)
  523. para_blocks.append(new_block)
  524. blocks["para_blocks"] = para_blocks
  525. return pdf_dic
  526. class BlockStatisticsCalculator:
  527. """
  528. This class calculates the statistics of the block.
  529. """
  530. def __init__(self) -> None:
  531. pass
  532. def __calc_stats_of_new_lines(self, new_lines):
  533. """
  534. This function calculates the paragraph metrics
  535. Parameters
  536. ----------
  537. combined_lines : list
  538. combined lines
  539. Returns
  540. -------
  541. X0 : float
  542. Median of x0 values, which represents the left average boundary of the block
  543. X1 : float
  544. Median of x1 values, which represents the right average boundary of the block
  545. avg_char_width : float
  546. Average of char widths, which represents the average char width of the block
  547. avg_char_height : float
  548. Average of line heights, which represents the average line height of the block
  549. """
  550. x0_values = []
  551. x1_values = []
  552. char_widths = []
  553. char_heights = []
  554. block_font_types = []
  555. block_font_sizes = []
  556. block_directions = []
  557. if len(new_lines) > 0:
  558. for i, line in enumerate(new_lines):
  559. line_bbox = line["bbox"]
  560. line_text = line["text"]
  561. line_spans = line["spans"]
  562. num_chars = len([ch for ch in line_text if not ch.isspace()])
  563. x0_values.append(line_bbox[0])
  564. x1_values.append(line_bbox[2])
  565. if num_chars > 0:
  566. char_width = (line_bbox[2] - line_bbox[0]) / num_chars
  567. char_widths.append(char_width)
  568. for span in line_spans:
  569. block_font_types.append(span["font"])
  570. block_font_sizes.append(span["size"])
  571. if "dir" in line:
  572. block_directions.append(line["dir"])
  573. # line_font_types = [span["font"] for span in line_spans]
  574. char_heights = [span["size"] for span in line_spans]
  575. X0 = np.median(x0_values) if x0_values else 0
  576. X1 = np.median(x1_values) if x1_values else 0
  577. avg_char_width = sum(char_widths) / len(char_widths) if char_widths else 0
  578. avg_char_height = sum(char_heights) / len(char_heights) if char_heights else 0
  579. # max_freq_font_type = max(set(block_font_types), key=block_font_types.count) if block_font_types else None
  580. max_span_length = 0
  581. max_span_font_type = None
  582. for line in new_lines:
  583. line_spans = line["spans"]
  584. for span in line_spans:
  585. span_length = span["bbox"][2] - span["bbox"][0]
  586. if span_length > max_span_length:
  587. max_span_length = span_length
  588. max_span_font_type = span["font"]
  589. max_freq_font_type = max_span_font_type
  590. avg_font_size = sum(block_font_sizes) / len(block_font_sizes) if block_font_sizes else None
  591. avg_dir_horizontal = sum([dir[0] for dir in block_directions]) / len(block_directions) if block_directions else 0
  592. avg_dir_vertical = sum([dir[1] for dir in block_directions]) / len(block_directions) if block_directions else 0
  593. median_font_size = float(np.median(block_font_sizes)) if block_font_sizes else None
  594. return (
  595. X0,
  596. X1,
  597. avg_char_width,
  598. avg_char_height,
  599. max_freq_font_type,
  600. avg_font_size,
  601. (avg_dir_horizontal, avg_dir_vertical),
  602. median_font_size,
  603. )
  604. def __make_new_block(self, input_block):
  605. new_block = {}
  606. raw_lines = input_block["lines"]
  607. stats = self.__calc_stats_of_new_lines(raw_lines)
  608. block_id = input_block["block_id"]
  609. block_bbox = input_block["bbox"]
  610. block_text = input_block["text"]
  611. block_lines = raw_lines
  612. block_avg_left_boundary = stats[0]
  613. block_avg_right_boundary = stats[1]
  614. block_avg_char_width = stats[2]
  615. block_avg_char_height = stats[3]
  616. block_font_type = stats[4]
  617. block_font_size = stats[5]
  618. block_direction = stats[6]
  619. block_median_font_size = stats[7]
  620. new_block["block_id"] = block_id
  621. new_block["bbox"] = block_bbox
  622. new_block["text"] = block_text
  623. new_block["dir"] = block_direction
  624. new_block["X0"] = block_avg_left_boundary
  625. new_block["X1"] = block_avg_right_boundary
  626. new_block["avg_char_width"] = block_avg_char_width
  627. new_block["avg_char_height"] = block_avg_char_height
  628. new_block["block_font_type"] = block_font_type
  629. new_block["block_font_size"] = block_font_size
  630. new_block["lines"] = block_lines
  631. new_block["median_font_size"] = block_median_font_size
  632. return new_block
  633. def batch_process_blocks(self, pdf_dic):
  634. """
  635. This function processes the blocks in batch.
  636. Parameters
  637. ----------
  638. self : object
  639. The instance of the class.
  640. ----------
  641. blocks : list
  642. Input block is a list of raw blocks.
  643. Schema can refer to the value of key ""preproc_blocks".
  644. Returns
  645. -------
  646. result_dict : dict
  647. result dictionary
  648. """
  649. for page_id, blocks in pdf_dic.items():
  650. if page_id.startswith("page_"):
  651. para_blocks = []
  652. if "para_blocks" in blocks.keys():
  653. input_blocks = blocks["para_blocks"]
  654. for input_block in input_blocks:
  655. new_block = self.__make_new_block(input_block)
  656. para_blocks.append(new_block)
  657. blocks["para_blocks"] = para_blocks
  658. return pdf_dic
  659. class DocStatisticsCalculator:
  660. """
  661. This class calculates the statistics of the document.
  662. """
  663. def __init__(self) -> None:
  664. pass
  665. def calc_stats_of_doc(self, pdf_dict):
  666. """
  667. This function computes the statistics of the document
  668. Parameters
  669. ----------
  670. result_dict : dict
  671. result dictionary
  672. Returns
  673. -------
  674. statistics : dict
  675. statistics of the document
  676. """
  677. total_text_length = 0
  678. total_num_blocks = 0
  679. for page_id, blocks in pdf_dict.items():
  680. if page_id.startswith("page_"):
  681. if "para_blocks" in blocks.keys():
  682. para_blocks = blocks["para_blocks"]
  683. for para_block in para_blocks:
  684. total_text_length += len(para_block["text"])
  685. total_num_blocks += 1
  686. avg_text_length = total_text_length / total_num_blocks if total_num_blocks else 0
  687. font_list = []
  688. for page_id, blocks in pdf_dict.items():
  689. if page_id.startswith("page_"):
  690. if "para_blocks" in blocks.keys():
  691. input_blocks = blocks["para_blocks"]
  692. for input_block in input_blocks:
  693. block_text_length = len(input_block.get("text", ""))
  694. if block_text_length < avg_text_length * 0.5:
  695. continue
  696. block_font_type = safe_get(input_block, "block_font_type", "")
  697. block_font_size = safe_get(input_block, "block_font_size", 0)
  698. font_list.append((block_font_type, block_font_size))
  699. font_counter = Counter(font_list)
  700. most_common_font = font_counter.most_common(1)[0] if font_list else (("", 0), 0)
  701. second_most_common_font = font_counter.most_common(2)[1] if len(font_counter) > 1 else (("", 0), 0)
  702. statistics = {
  703. "num_pages": 0,
  704. "num_blocks": 0,
  705. "num_paras": 0,
  706. "num_titles": 0,
  707. "num_header_blocks": 0,
  708. "num_footer_blocks": 0,
  709. "num_watermark_blocks": 0,
  710. "num_vertical_margin_note_blocks": 0,
  711. "most_common_font_type": most_common_font[0][0],
  712. "most_common_font_size": most_common_font[0][1],
  713. "number_of_most_common_font": most_common_font[1],
  714. "second_most_common_font_type": second_most_common_font[0][0],
  715. "second_most_common_font_size": second_most_common_font[0][1],
  716. "number_of_second_most_common_font": second_most_common_font[1],
  717. "avg_text_length": avg_text_length,
  718. }
  719. for page_id, blocks in pdf_dict.items():
  720. if page_id.startswith("page_"):
  721. blocks = pdf_dict[page_id]["para_blocks"]
  722. statistics["num_pages"] += 1
  723. for block_id, block_data in enumerate(blocks):
  724. statistics["num_blocks"] += 1
  725. if "paras" in block_data.keys():
  726. statistics["num_paras"] += len(block_data["paras"])
  727. for line in block_data["lines"]:
  728. if line.get("is_title", 0):
  729. statistics["num_titles"] += 1
  730. if block_data.get("is_header", 0):
  731. statistics["num_header_blocks"] += 1
  732. if block_data.get("is_footer", 0):
  733. statistics["num_footer_blocks"] += 1
  734. if block_data.get("is_watermark", 0):
  735. statistics["num_watermark_blocks"] += 1
  736. if block_data.get("is_vertical_margin_note", 0):
  737. statistics["num_vertical_margin_note_blocks"] += 1
  738. pdf_dict["statistics"] = statistics
  739. return pdf_dict
  740. class TitleProcessor:
  741. """
  742. This class processes the title.
  743. """
  744. def __init__(self, *doc_statistics) -> None:
  745. if len(doc_statistics) > 0:
  746. self.doc_statistics = doc_statistics[0]
  747. self.nlp_model = NLPModels()
  748. self.MAX_TITLE_LEVEL = 3
  749. self.numbered_title_pattern = r"""
  750. ^ # 行首
  751. ( # 开始捕获组
  752. [\(\(]\d+[\)\)] # 括号内数字,支持中文和英文括号,例如:(1) 或 (1)
  753. |\d+[\)\)]\s # 数字后跟右括号和空格,支持中文和英文括号,例如:2) 或 2)
  754. |[\(\(][A-Z][\)\)] # 括号内大写字母,支持中文和英文括号,例如:(A) 或 (A)
  755. |[A-Z][\)\)]\s # 大写字母后跟右括号和空格,例如:A) 或 A)
  756. |[\(\(][IVXLCDM]+[\)\)] # 括号内罗马数字,支持中文和英文括号,例如:(I) 或 (I)
  757. |[IVXLCDM]+[\)\)]\s # 罗马数字后跟右括号和空格,例如:I) 或 I)
  758. |\d+(\.\d+)*\s # 数字或复合数字编号后跟空格,例如:1. 或 3.2.1
  759. |[一二三四五六七八九十百千]+[、\s] # 中文序号后跟顿号和空格,例如:一、
  760. |[\(|\(][一二三四五六七八九十百千]+[\)|\)]\s* # 中文括号内中文序号后跟空格,例如:(一)
  761. |[A-Z]\.\d+(\.\d+)?\s # 大写字母后跟点和数字,例如:A.1 或 A.1.1
  762. |[\(\(][a-z][\)\)] # 括号内小写字母,支持中文和英文括号,例如:(a) 或 (a)
  763. |[a-z]\)\s # 小写字母后跟右括号和空格,例如:a)
  764. |[A-Z]-\s # 大写字母后跟短横线和空格,例如:A-
  765. |\w+:\s # 英文序号词后跟冒号和空格,例如:First:
  766. |第[一二三四五六七八九十百千]+[章节部分条款]\s # 以“第”开头的中文标题后跟空格
  767. |[IVXLCDM]+\. # 罗马数字后跟点,例如:I.
  768. |\d+\.\s # 单个数字后跟点和空格,例如:1.
  769. ) # 结束捕获组
  770. .+ # 标题的其余部分
  771. """
  772. self.stage = (
  773. 0 # Used for distinguishing the stage of title detection, the number is occurred in paragraph process pipeline
  774. )
  775. def _is_potential_title(
  776. self,
  777. curr_line,
  778. prev_line,
  779. prev_line_is_title,
  780. next_line,
  781. avg_char_width,
  782. avg_char_height,
  783. median_font_size,
  784. ):
  785. """
  786. This function checks if the line is a potential title.
  787. Parameters
  788. ----------
  789. curr_line : dict
  790. current line
  791. prev_line : dict
  792. previous line
  793. next_line : dict
  794. next line
  795. avg_char_width : float
  796. average of char widths
  797. avg_char_height : float
  798. average of line heights
  799. Returns
  800. -------
  801. bool
  802. True if the line is a potential title, False otherwise.
  803. """
  804. def __is_line_centered(line_bbox, page_bbox, avg_char_width):
  805. """
  806. This function checks if the line is centered on the page
  807. Parameters
  808. ----------
  809. line_bbox : list
  810. bbox of the line
  811. page_bbox : list
  812. bbox of the page
  813. avg_char_width : float
  814. average of char widths
  815. Returns
  816. -------
  817. bool
  818. True if the line is centered on the page, False otherwise.
  819. """
  820. horizontal_ratio = 0.5
  821. horizontal_thres = horizontal_ratio * avg_char_width
  822. x0, _, x1, _ = line_bbox
  823. _, _, page_x1, _ = page_bbox
  824. return abs((x0 + x1) / 2 - page_x1 / 2) < horizontal_thres
  825. def __is_bold_font_line(line):
  826. """
  827. Check if a line contains any bold font style.
  828. """
  829. def _is_bold_span(span):
  830. # if span text is empty or only contains space, return False
  831. if not span["text"].strip():
  832. return False
  833. return bool(span["flags"] & 2**4) # Check if the font is bold
  834. for span in line["spans"]:
  835. if not _is_bold_span(span):
  836. return False
  837. return True
  838. def __is_italic_font_line(line):
  839. """
  840. Check if a line contains any italic font style.
  841. """
  842. def __is_italic_span(span):
  843. return bool(span["flags"] & 2**1) # Check if the font is italic
  844. for span in line["spans"]:
  845. if not __is_italic_span(span):
  846. return False
  847. return True
  848. def __is_punctuation_heavy(line_text):
  849. """
  850. Check if the line contains a high ratio of punctuation marks, which may indicate
  851. that the line is not a title.
  852. Parameters:
  853. line_text (str): Text of the line.
  854. Returns:
  855. bool: True if the line is heavy with punctuation, False otherwise.
  856. """
  857. # Pattern for common title format like "X.Y. Title"
  858. pattern = r"\b\d+\.\d+\..*\b"
  859. # If the line matches the title format, return False
  860. if re.match(pattern, line_text.strip()):
  861. return False
  862. # Find all punctuation marks in the line
  863. punctuation_marks = re.findall(r"[^\w\s]", line_text)
  864. number_of_punctuation_marks = len(punctuation_marks)
  865. text_length = len(line_text)
  866. if text_length == 0:
  867. return False
  868. punctuation_ratio = number_of_punctuation_marks / text_length
  869. if punctuation_ratio >= 0.1:
  870. return True
  871. return False
  872. def __has_mixed_font_styles(spans, strict_mode=False):
  873. """
  874. This function checks if the line has mixed font styles, the strict mode will compare the font types
  875. Parameters
  876. ----------
  877. spans : list
  878. spans of the line
  879. strict_mode : bool
  880. True for strict mode, the font types will be fully compared
  881. False for non-strict mode, the font types will be compared by the most longest common prefix
  882. Returns
  883. -------
  884. bool
  885. True if the line has mixed font styles, False otherwise.
  886. """
  887. if strict_mode:
  888. font_styles = set()
  889. for span in spans:
  890. font_style = span["font"].lower()
  891. font_styles.add(font_style)
  892. return len(font_styles) > 1
  893. else: # non-strict mode
  894. font_styles = []
  895. for span in spans:
  896. font_style = span["font"].lower()
  897. font_styles.append(font_style)
  898. if len(font_styles) > 1:
  899. longest_common_prefix = os.path.commonprefix(font_styles)
  900. if len(longest_common_prefix) > 0:
  901. return False
  902. else:
  903. return True
  904. else:
  905. return False
  906. def __is_different_font_type_from_neighbors(curr_line_font_type, prev_line_font_type, next_line_font_type):
  907. """
  908. This function checks if the current line has a different font type from the previous and next lines
  909. Parameters
  910. ----------
  911. curr_line_font_type : str
  912. font type of the current line
  913. prev_line_font_type : str
  914. font type of the previous line
  915. next_line_font_type : str
  916. font type of the next line
  917. Returns
  918. -------
  919. bool
  920. True if the current line has a different font type from the previous and next lines, False otherwise.
  921. """
  922. return all(
  923. curr_line_font_type != other_font_type.lower()
  924. for other_font_type in [prev_line_font_type, next_line_font_type]
  925. if other_font_type is not None
  926. )
  927. def __is_larger_font_size_from_neighbors(curr_line_font_size, prev_line_font_size, next_line_font_size):
  928. """
  929. This function checks if the current line has a larger font size than the previous and next lines
  930. Parameters
  931. ----------
  932. curr_line_font_size : float
  933. font size of the current line
  934. prev_line_font_size : float
  935. font size of the previous line
  936. next_line_font_size : float
  937. font size of the next line
  938. Returns
  939. -------
  940. bool
  941. True if the current line has a larger font size than the previous and next lines, False otherwise.
  942. """
  943. return all(
  944. curr_line_font_size > other_font_size * 1.2
  945. for other_font_size in [prev_line_font_size, next_line_font_size]
  946. if other_font_size is not None
  947. )
  948. def __is_similar_to_pre_line(curr_line_font_type, prev_line_font_type, curr_line_font_size, prev_line_font_size):
  949. """
  950. This function checks if the current line is similar to the previous line
  951. Parameters
  952. ----------
  953. curr_line : dict
  954. current line
  955. prev_line : dict
  956. previous line
  957. Returns
  958. -------
  959. bool
  960. True if the current line is similar to the previous line, False otherwise.
  961. """
  962. if curr_line_font_type == prev_line_font_type and curr_line_font_size == prev_line_font_size:
  963. return True
  964. else:
  965. return False
  966. def __is_same_font_type_of_docAvg(curr_line_font_type):
  967. """
  968. This function checks if the current line has the same font type as the document average font type
  969. Parameters
  970. ----------
  971. curr_line_font_type : str
  972. font type of the current line
  973. Returns
  974. -------
  975. bool
  976. True if the current line has the same font type as the document average font type, False otherwise.
  977. """
  978. doc_most_common_font_type = safe_get(self.doc_statistics, "most_common_font_type", "").lower()
  979. doc_second_most_common_font_type = safe_get(self.doc_statistics, "second_most_common_font_type", "").lower()
  980. return curr_line_font_type.lower() in [doc_most_common_font_type, doc_second_most_common_font_type]
  981. def __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio: float = 1):
  982. """
  983. This function checks if the current line has a large enough font size
  984. Parameters
  985. ----------
  986. curr_line_font_size : float
  987. font size of the current line
  988. ratio : float
  989. ratio of the current line font size to the document average font size
  990. Returns
  991. -------
  992. bool
  993. True if the current line has a large enough font size, False otherwise.
  994. """
  995. doc_most_common_font_size = safe_get(self.doc_statistics, "most_common_font_size", 0)
  996. doc_second_most_common_font_size = safe_get(self.doc_statistics, "second_most_common_font_size", 0)
  997. doc_avg_font_size = min(doc_most_common_font_size, doc_second_most_common_font_size)
  998. return curr_line_font_size >= doc_avg_font_size * ratio
  999. def __is_sufficient_spacing_above_and_below(
  1000. curr_line_bbox,
  1001. prev_line_bbox,
  1002. next_line_bbox,
  1003. avg_char_height,
  1004. median_font_size,
  1005. ):
  1006. """
  1007. This function checks if the current line has sufficient spacing above and below
  1008. Parameters
  1009. ----------
  1010. curr_line_bbox : list
  1011. bbox of the current line
  1012. prev_line_bbox : list
  1013. bbox of the previous line
  1014. next_line_bbox : list
  1015. bbox of the next line
  1016. avg_char_width : float
  1017. average of char widths
  1018. avg_char_height : float
  1019. average of line heights
  1020. Returns
  1021. -------
  1022. bool
  1023. True if the current line has sufficient spacing above and below, False otherwise.
  1024. """
  1025. vertical_ratio = 1.25
  1026. vertical_thres = vertical_ratio * median_font_size
  1027. _, y0, _, y1 = curr_line_bbox
  1028. sufficient_spacing_above = False
  1029. if prev_line_bbox:
  1030. vertical_spacing_above = min(y0 - prev_line_bbox[1], y1 - prev_line_bbox[3])
  1031. sufficient_spacing_above = vertical_spacing_above > vertical_thres
  1032. else:
  1033. sufficient_spacing_above = True
  1034. sufficient_spacing_below = False
  1035. if next_line_bbox:
  1036. vertical_spacing_below = min(next_line_bbox[1] - y0, next_line_bbox[3] - y1)
  1037. sufficient_spacing_below = vertical_spacing_below > vertical_thres
  1038. else:
  1039. sufficient_spacing_below = True
  1040. return (sufficient_spacing_above, sufficient_spacing_below)
  1041. def __is_word_list_line_by_rules(curr_line_text):
  1042. """
  1043. This function checks if the current line is a word list
  1044. Parameters
  1045. ----------
  1046. curr_line_text : str
  1047. text of the current line
  1048. Returns
  1049. -------
  1050. bool
  1051. True if the current line is a name list, False otherwise.
  1052. """
  1053. # name_list_pattern = r"([a-zA-Z][a-zA-Z\s]{0,20}[a-zA-Z]|[\u4e00-\u9fa5·]{2,16})(?=[,,;;\s]|$)"
  1054. name_list_pattern = r"(?<![\u4e00-\u9fa5])([A-Z][a-z]{0,19}\s[A-Z][a-z]{0,19}|[\u4e00-\u9fa5]{2,6})(?=[,,;;\s]|$)"
  1055. compiled_pattern = re.compile(name_list_pattern)
  1056. if compiled_pattern.search(curr_line_text):
  1057. return True
  1058. else:
  1059. return False
  1060. def __get_text_catgr_by_nlp(curr_line_text):
  1061. """
  1062. This function checks if the current line is a name list using nlp model, such as spacy
  1063. Parameters
  1064. ----------
  1065. curr_line_text : str
  1066. text of the current line
  1067. Returns
  1068. -------
  1069. bool
  1070. True if the current line is a name list, False otherwise.
  1071. """
  1072. result = self.nlp_model.detect_entity_catgr_using_nlp(curr_line_text)
  1073. return result
  1074. def __is_numbered_title(curr_line_text):
  1075. """
  1076. This function checks if the current line is a numbered list
  1077. Parameters
  1078. ----------
  1079. curr_line_text : str
  1080. text of the current line
  1081. Returns
  1082. -------
  1083. bool
  1084. True if the current line is a numbered list, False otherwise.
  1085. """
  1086. compiled_pattern = re.compile(self.numbered_title_pattern, re.VERBOSE)
  1087. if compiled_pattern.search(curr_line_text):
  1088. return True
  1089. else:
  1090. return False
  1091. def __is_end_with_ending_puncs(line_text):
  1092. """
  1093. This function checks if the current line ends with a ending punctuation mark
  1094. Parameters
  1095. ----------
  1096. line_text : str
  1097. text of the current line
  1098. Returns
  1099. -------
  1100. bool
  1101. True if the current line ends with a punctuation mark, False otherwise.
  1102. """
  1103. end_puncs = [".", "?", "!", "。", "?", "!", "…"]
  1104. line_text = line_text.rstrip()
  1105. if line_text[-1] in end_puncs:
  1106. return True
  1107. return False
  1108. def __contains_only_no_meaning_symbols(line_text):
  1109. """
  1110. This function checks if the current line contains only symbols that have no meaning, if so, it is not a title.
  1111. Situation contains:
  1112. 1. Only have punctuation marks
  1113. 2. Only have other non-meaning symbols
  1114. Parameters
  1115. ----------
  1116. line_text : str
  1117. text of the current line
  1118. Returns
  1119. -------
  1120. bool
  1121. True if the current line contains only symbols that have no meaning, False otherwise.
  1122. """
  1123. punctuation_marks = re.findall(r"[^\w\s]", line_text) # find all punctuation marks
  1124. number_of_punctuation_marks = len(punctuation_marks)
  1125. text_length = len(line_text)
  1126. if text_length == 0:
  1127. return False
  1128. punctuation_ratio = number_of_punctuation_marks / text_length
  1129. if punctuation_ratio >= 0.9:
  1130. return True
  1131. return False
  1132. def __is_equation(line_text):
  1133. """
  1134. This function checks if the current line is an equation.
  1135. Parameters
  1136. ----------
  1137. line_text : str
  1138. Returns
  1139. -------
  1140. bool
  1141. True if the current line is an equation, False otherwise.
  1142. """
  1143. equation_reg = r"\$.*?\\overline.*?\$" # to match interline equations
  1144. if re.search(equation_reg, line_text):
  1145. return True
  1146. else:
  1147. return False
  1148. def __is_title_by_len(text, max_length=200):
  1149. """
  1150. This function checks if the current line is a title by length.
  1151. Parameters
  1152. ----------
  1153. text : str
  1154. text of the current line
  1155. max_length : int
  1156. max length of the title
  1157. Returns
  1158. -------
  1159. bool
  1160. True if the current line is a title, False otherwise.
  1161. """
  1162. text = text.strip()
  1163. return len(text) <= max_length
  1164. def __compute_line_font_type_and_size(curr_line):
  1165. """
  1166. This function computes the font type and font size of the line.
  1167. Parameters
  1168. ----------
  1169. line : dict
  1170. line
  1171. Returns
  1172. -------
  1173. font_type : str
  1174. font type of the line
  1175. font_size : float
  1176. font size of the line
  1177. """
  1178. spans = curr_line["spans"]
  1179. max_accumulated_length = 0
  1180. max_span_font_size = curr_line["spans"][0]["size"] # default value, float type
  1181. max_span_font_type = curr_line["spans"][0]["font"].lower() # default value, string type
  1182. for span in spans:
  1183. if span["text"].isspace():
  1184. continue
  1185. span_length = span["bbox"][2] - span["bbox"][0]
  1186. if span_length > max_accumulated_length:
  1187. max_accumulated_length = span_length
  1188. max_span_font_size = span["size"]
  1189. max_span_font_type = span["font"].lower()
  1190. return max_span_font_type, max_span_font_size
  1191. def __is_a_consistent_sub_title(pre_line, curr_line):
  1192. """
  1193. This function checks if the current line is a consistent sub title.
  1194. Parameters
  1195. ----------
  1196. pre_line : dict
  1197. previous line
  1198. curr_line : dict
  1199. current line
  1200. Returns
  1201. -------
  1202. bool
  1203. True if the current line is a consistent sub title, False otherwise.
  1204. """
  1205. if pre_line is None:
  1206. return False
  1207. start_letter_of_pre_line = pre_line["text"][0]
  1208. start_letter_of_curr_line = curr_line["text"][0]
  1209. has_same_prefix_digit = (
  1210. start_letter_of_pre_line.isdigit()
  1211. and start_letter_of_curr_line.isdigit()
  1212. and start_letter_of_pre_line == start_letter_of_curr_line
  1213. )
  1214. # prefix text of curr_line satisfies the following title format: x.x
  1215. prefix_text_pattern = r"^\d+\.\d+"
  1216. subtitle_format_match = re.match(prefix_text_pattern, curr_line["text"])
  1217. if subtitle_format_match:
  1218. has_subtitle_format = True
  1219. else:
  1220. has_subtitle_format = False
  1221. if has_same_prefix_digit or has_subtitle_format:
  1222. print("is a consistent sub title")
  1223. return True
  1224. """
  1225. Title detecting main Process.
  1226. """
  1227. """
  1228. Basic features about the current line.
  1229. """
  1230. curr_line_bbox = curr_line["bbox"]
  1231. curr_line_text = curr_line["text"]
  1232. curr_line_font_type, curr_line_font_size = __compute_line_font_type_and_size(curr_line)
  1233. if len(curr_line_text.strip()) == 0: # skip empty lines
  1234. return False, False
  1235. prev_line_bbox = prev_line["bbox"] if prev_line else None
  1236. if prev_line:
  1237. prev_line_font_type, prev_line_font_size = __compute_line_font_type_and_size(prev_line)
  1238. else:
  1239. prev_line_font_type, prev_line_font_size = None, None
  1240. next_line_bbox = next_line["bbox"] if next_line else None
  1241. if next_line:
  1242. next_line_font_type, next_line_font_size = __compute_line_font_type_and_size(next_line)
  1243. else:
  1244. next_line_font_type, next_line_font_size = None, None
  1245. """
  1246. Aggregated features about the current line.
  1247. """
  1248. is_italc_font = __is_italic_font_line(curr_line)
  1249. is_bold_font = __is_bold_font_line(curr_line)
  1250. is_font_size_little_less_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=0.8)
  1251. is_font_size_not_less_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=1)
  1252. is_much_larger_font_than_doc_avg = __is_font_size_not_less_than_docAvg(curr_line_font_size, ratio=1.6)
  1253. is_not_same_font_type_of_docAvg = not __is_same_font_type_of_docAvg(curr_line_font_type)
  1254. is_potential_title_font = is_bold_font or is_font_size_not_less_than_doc_avg or is_not_same_font_type_of_docAvg
  1255. is_mix_font_styles_strict = __has_mixed_font_styles(curr_line["spans"], strict_mode=True)
  1256. is_mix_font_styles_loose = __has_mixed_font_styles(curr_line["spans"], strict_mode=False)
  1257. is_punctuation_heavy = __is_punctuation_heavy(curr_line_text)
  1258. is_word_list_line_by_rules = __is_word_list_line_by_rules(curr_line_text)
  1259. is_person_or_org_list_line_by_nlp = __get_text_catgr_by_nlp(curr_line_text) in ["PERSON", "GPE", "ORG"]
  1260. is_font_size_larger_than_neighbors = __is_larger_font_size_from_neighbors(
  1261. curr_line_font_size, prev_line_font_size, next_line_font_size
  1262. )
  1263. is_font_type_diff_from_neighbors = __is_different_font_type_from_neighbors(
  1264. curr_line_font_type, prev_line_font_type, next_line_font_type
  1265. )
  1266. has_sufficient_spaces_above, has_sufficient_spaces_below = __is_sufficient_spacing_above_and_below(
  1267. curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_height, median_font_size
  1268. )
  1269. is_similar_to_pre_line = __is_similar_to_pre_line(
  1270. curr_line_font_type, prev_line_font_type, curr_line_font_size, prev_line_font_size
  1271. )
  1272. is_consistent_sub_title = __is_a_consistent_sub_title(prev_line, curr_line)
  1273. """
  1274. Further aggregated features about the current line.
  1275. Attention:
  1276. Features that start with __ are for internal use.
  1277. """
  1278. __is_line_left_aligned_from_neighbors = is_line_left_aligned_from_neighbors(
  1279. curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width
  1280. )
  1281. __is_font_diff_from_neighbors = is_font_size_larger_than_neighbors or is_font_type_diff_from_neighbors
  1282. is_a_left_inline_title = (
  1283. is_mix_font_styles_strict and __is_line_left_aligned_from_neighbors and __is_font_diff_from_neighbors
  1284. )
  1285. is_title_by_check_prev_line = prev_line is None and has_sufficient_spaces_above and is_potential_title_font
  1286. is_title_by_check_next_line = next_line is None and has_sufficient_spaces_below and is_potential_title_font
  1287. is_title_by_check_pre_and_next_line = (
  1288. (prev_line is not None or next_line is not None)
  1289. and has_sufficient_spaces_above
  1290. and has_sufficient_spaces_below
  1291. and is_potential_title_font
  1292. )
  1293. is_numbered_title = __is_numbered_title(curr_line_text) and (
  1294. (has_sufficient_spaces_above or prev_line is None) and (has_sufficient_spaces_below or next_line is None)
  1295. )
  1296. is_not_end_with_ending_puncs = not __is_end_with_ending_puncs(curr_line_text)
  1297. is_not_only_no_meaning_symbols = not __contains_only_no_meaning_symbols(curr_line_text)
  1298. is_equation = __is_equation(curr_line_text)
  1299. is_title_by_len = __is_title_by_len(curr_line_text)
  1300. """
  1301. Decide if the line is a title.
  1302. """
  1303. is_title = (
  1304. is_not_end_with_ending_puncs # not end with ending punctuation marks
  1305. and is_not_only_no_meaning_symbols # not only have no meaning symbols
  1306. and is_title_by_len # is a title by length, default max length is 200
  1307. and not is_equation # an interline equation should never be a title
  1308. and is_potential_title_font # is a potential title font, which is bold or larger than the document average font size or not the same font type as the document average font type
  1309. and (
  1310. (is_not_same_font_type_of_docAvg and is_font_size_not_less_than_doc_avg)
  1311. or (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
  1312. or (
  1313. is_much_larger_font_than_doc_avg
  1314. and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
  1315. )
  1316. or (
  1317. is_font_size_little_less_than_doc_avg
  1318. and is_bold_font
  1319. and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
  1320. )
  1321. ) # not the same font type as the document average font type, which includes the most common font type and the second most common font type
  1322. and (
  1323. (
  1324. not is_person_or_org_list_line_by_nlp
  1325. and (
  1326. is_much_larger_font_than_doc_avg
  1327. or (is_not_same_font_type_of_docAvg and is_font_size_not_less_than_doc_avg)
  1328. )
  1329. )
  1330. or (
  1331. not (is_word_list_line_by_rules and is_person_or_org_list_line_by_nlp)
  1332. and not is_a_left_inline_title
  1333. and not is_punctuation_heavy
  1334. and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
  1335. )
  1336. or (
  1337. is_person_or_org_list_line_by_nlp
  1338. and (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
  1339. and (is_bold_font and is_much_larger_font_than_doc_avg and is_not_same_font_type_of_docAvg)
  1340. )
  1341. or (is_numbered_title and not is_a_left_inline_title)
  1342. )
  1343. # )
  1344. ) or (prev_line_is_title and is_consistent_sub_title)
  1345. is_name_or_org_list_to_be_removed = (
  1346. (is_person_or_org_list_line_by_nlp)
  1347. and is_punctuation_heavy
  1348. and (is_title_by_check_prev_line or is_title_by_check_next_line or is_title_by_check_pre_and_next_line)
  1349. ) and not is_title
  1350. if is_name_or_org_list_to_be_removed:
  1351. is_author_or_org_list = True
  1352. else:
  1353. is_author_or_org_list = False
  1354. # return is_title, is_author_or_org_list
  1355. # """
  1356. """
  1357. # print reason why the line is a title
  1358. if is_title:
  1359. print_green("This line is a title.")
  1360. print_green("↓" * 10)
  1361. print()
  1362. print("curr_line_text: ", curr_line_text)
  1363. print()
  1364. print(f"prev_line_is_title: {prev_line_is_title}")
  1365. print()
  1366. print(f"is_consistent_sub_title: {is_consistent_sub_title}")
  1367. """
  1368. # print reason why the line is not a title
  1369. # line_text = curr_line_text.strip()
  1370. # test_text = "Career/Personal Life"
  1371. # text_content_condition = line_text == test_text
  1372. # if not is_title and text_content_condition: # Print specific line
  1373. """
  1374. if not is_title: # Print each line
  1375. print_red("This line is not a title.")
  1376. print_red("↓" * 10)
  1377. print()
  1378. print("curr_line_text: ", curr_line_text)
  1379. print()
  1380. if is_not_end_with_ending_puncs:
  1381. print_green(f"is_not_end_with_ending_puncs")
  1382. else:
  1383. print_red(f"is_end_with_ending_puncs")
  1384. if is_not_only_no_meaning_symbols:
  1385. print_green(f"is_not_only_no_meaning_symbols")
  1386. else:
  1387. print_red(f"is_only_no_meaning_symbols")
  1388. if is_title_by_len:
  1389. print_green(f"is_title_by_len: {is_title_by_len}")
  1390. else:
  1391. print_red(f"is_not_title_by_len: {is_title_by_len}")
  1392. if is_equation:
  1393. print_red(f"is_equation")
  1394. else:
  1395. print_green(f"is_not_equation")
  1396. if is_potential_title_font:
  1397. print_green(f"is_potential_title_font")
  1398. else:
  1399. print_red(f"is_not_potential_title_font")
  1400. if is_punctuation_heavy:
  1401. print_red("is_punctuation_heavy")
  1402. else:
  1403. print_green("is_not_punctuation_heavy")
  1404. if is_bold_font:
  1405. print_green(f"is_bold_font")
  1406. else:
  1407. print_red(f"is_not_bold_font")
  1408. if is_font_size_not_less_than_doc_avg:
  1409. print_green(f"is_larger_font_than_doc_avg")
  1410. else:
  1411. print_red(f"is_not_larger_font_than_doc_avg")
  1412. if is_much_larger_font_than_doc_avg:
  1413. print_green(f"is_much_larger_font_than_doc_avg")
  1414. else:
  1415. print_red(f"is_not_much_larger_font_than_doc_avg")
  1416. if is_not_same_font_type_of_docAvg:
  1417. print_green(f"is_not_same_font_type_of_docAvg")
  1418. else:
  1419. print_red(f"is_same_font_type_of_docAvg")
  1420. if is_word_list_line_by_rules:
  1421. print_red("is_word_list_line_by_rules")
  1422. else:
  1423. print_green("is_not_name_list_by_rules")
  1424. if is_person_or_org_list_line_by_nlp:
  1425. print_red("is_person_or_org_list_line_by_nlp")
  1426. else:
  1427. print_green("is_not_person_or_org_list_line_by_nlp")
  1428. if not is_numbered_title:
  1429. print_red("is_not_numbered_title")
  1430. else:
  1431. print_green("is_numbered_title")
  1432. if is_a_left_inline_title:
  1433. print_red("is_a_left_inline_title")
  1434. else:
  1435. print_green("is_not_a_left_inline_title")
  1436. if not is_title_by_check_prev_line:
  1437. print_red("is_not_title_by_check_prev_line")
  1438. else:
  1439. print_green("is_title_by_check_prev_line")
  1440. if not is_title_by_check_next_line:
  1441. print_red("is_not_title_by_check_next_line")
  1442. else:
  1443. print_green("is_title_by_check_next_line")
  1444. if not is_title_by_check_pre_and_next_line:
  1445. print_red("is_not_title_by_check_pre_and_next_line")
  1446. else:
  1447. print_green("is_title_by_check_pre_and_next_line")
  1448. # print_green("Common features:")
  1449. # print_green("↓" * 10)
  1450. # print(f" curr_line_font_type: {curr_line_font_type}")
  1451. # print(f" curr_line_font_size: {curr_line_font_size}")
  1452. # print()
  1453. """
  1454. # """
  1455. return is_title, is_author_or_org_list
  1456. def _detect_title(self, curr_block, pre_block):
  1457. """
  1458. Use the functions 'is_potential_title' to detect titles of each paragraph block.
  1459. If a line is a title, then the value of key 'is_title' of the line will be set to True.
  1460. """
  1461. raw_lines = curr_block["lines"]
  1462. blk_avg_char_width = curr_block["avg_char_width"]
  1463. blk_avg_char_height = curr_block["avg_char_height"]
  1464. blk_media_font_size = curr_block["median_font_size"]
  1465. if self.stage == 0:
  1466. is_prev_line_a_title = False
  1467. for i, curr_line in enumerate(raw_lines):
  1468. prev_line = raw_lines[i - 1] if i > 0 else None
  1469. next_line = raw_lines[i + 1] if i < len(raw_lines) - 1 else None
  1470. is_line_a_title, is_line_an_entities_list = self._is_potential_title(
  1471. curr_line,
  1472. prev_line,
  1473. is_prev_line_a_title,
  1474. next_line,
  1475. blk_avg_char_width,
  1476. blk_avg_char_height,
  1477. blk_media_font_size,
  1478. )
  1479. if is_line_a_title:
  1480. curr_line["is_title"] = is_line_a_title
  1481. is_prev_line_a_title = True # set the flag to True for the next line
  1482. else:
  1483. curr_line["is_title"] = False
  1484. is_prev_line_a_title = False # set the flag to False for the next line
  1485. if is_line_an_entities_list:
  1486. curr_line["is_author_or_org_list"] = is_line_an_entities_list
  1487. else:
  1488. curr_line["is_author_or_org_list"] = False
  1489. return curr_block
  1490. if self.stage == 1: # Check the block consistent titles.
  1491. if pre_block and "paras" in pre_block.keys():
  1492. print_red(f"Checking cross block title...")
  1493. last_para_content = None
  1494. paras_of_pre_block = pre_block["paras"]
  1495. last_key = sorted(paras_of_pre_block.keys())[-1]
  1496. last_para_content = paras_of_pre_block[last_key]
  1497. if last_para_content is not None:
  1498. last_line_of_last_para_of_last_block = pre_block["lines"][-1]
  1499. first_line_of_curr_block = raw_lines[0]
  1500. next_line_of_curr_block = raw_lines[1] if len(raw_lines) > 1 else None
  1501. is_line_a_title, is_line_an_entities_list = self._is_potential_title(
  1502. first_line_of_curr_block,
  1503. last_line_of_last_para_of_last_block,
  1504. last_line_of_last_para_of_last_block["is_title"],
  1505. next_line_of_curr_block,
  1506. blk_avg_char_width,
  1507. blk_avg_char_height,
  1508. blk_media_font_size,
  1509. )
  1510. if is_line_a_title:
  1511. first_line_of_curr_block["is_title"] = is_line_a_title
  1512. else:
  1513. first_line_of_curr_block["is_title"] = False
  1514. if is_line_an_entities_list:
  1515. first_line_of_curr_block["is_author_or_org_list"] = is_line_an_entities_list
  1516. else:
  1517. first_line_of_curr_block["is_author_or_org_list"] = False
  1518. # print(f"first_line_of_curr_block: {first_line_of_curr_block['text']}")
  1519. # print(f"last_line_of_pre_block: {last_line_of_last_para['text']}")
  1520. return curr_block
  1521. else:
  1522. print_red(f"last_para_content is None")
  1523. def batch_detect_titles(self, pdf_dic):
  1524. """
  1525. This function batch process the blocks to detect titles.
  1526. Parameters
  1527. ----------
  1528. pdf_dict : dict
  1529. result dictionary
  1530. Returns
  1531. -------
  1532. pdf_dict : dict
  1533. result dictionary
  1534. """
  1535. num_titles = 0
  1536. for page_id, page_content in pdf_dic.items():
  1537. if page_id.startswith("page_"):
  1538. para_blocks = []
  1539. if "para_blocks" in page_content.keys():
  1540. para_blocks = page_content["para_blocks"]
  1541. all_single_line_blocks = []
  1542. for block in para_blocks:
  1543. if len(block["lines"]) == 1:
  1544. all_single_line_blocks.append(block)
  1545. new_para_blocks = []
  1546. if not len(all_single_line_blocks) == len(para_blocks): # Not all blocks are single line blocks.
  1547. for para_idx, para_block in enumerate(para_blocks):
  1548. print(f"______________________________________________________")
  1549. print(f"page_id: {page_id}")
  1550. print(f"para_block id: {para_block['block_id']}")
  1551. print(f"para_idx: {para_idx}")
  1552. pre_block = para_blocks[para_idx - 1] if para_idx > 0 else None
  1553. curr_block = para_block
  1554. print_yellow(f"text of current block: {curr_block['text'] if curr_block else None}")
  1555. print_green(f"text of previous block: {pre_block['text'] if pre_block else None}")
  1556. new_block = self._detect_title(curr_block, pre_block)
  1557. new_para_blocks.append(new_block)
  1558. # num_titles += sum([line.get("is_title", 0) for line in new_block["lines"]])
  1559. if new_block is not None:
  1560. num_titles += sum([line.get("is_title", 0) for line in new_block["lines"]])
  1561. else:
  1562. num_titles += 0
  1563. else: # All blocks are single line blocks.
  1564. for para_block in para_blocks:
  1565. new_para_blocks.append(para_block)
  1566. num_titles += sum([line.get("is_title", 0) for line in para_block["lines"]])
  1567. para_blocks = new_para_blocks
  1568. page_content["para_blocks"] = para_blocks
  1569. for para_block in para_blocks:
  1570. if para_block is not None:
  1571. all_titles = all(safe_get(line, "is_title", False) for line in para_block["lines"])
  1572. para_text_len = sum([len(line["text"]) for line in para_block["lines"]])
  1573. if (
  1574. all_titles and para_text_len < 200
  1575. ): # total length of the paragraph is less than 200, more than this should not be a title
  1576. para_block["is_block_title"] = 1
  1577. else:
  1578. para_block["is_block_title"] = 0
  1579. all_name_or_org_list_to_be_removed = all(
  1580. safe_get(line, "is_author_or_org_list", False) for line in para_block["lines"]
  1581. )
  1582. if all_name_or_org_list_to_be_removed and page_id == "page_0":
  1583. para_block["is_block_an_author_or_org_list"] = 1
  1584. else:
  1585. para_block["is_block_an_author_or_org_list"] = 0
  1586. else:
  1587. all_titles = False
  1588. # para_block["is_block_title"] = 0
  1589. # para_block["is_block_an_author_or_org_list"] = 0
  1590. # page_content["para_blocks"] = para_blocks
  1591. pdf_dic["statistics"]["num_titles"] = num_titles
  1592. return pdf_dic
  1593. def _recog_title_level(self, title_blocks):
  1594. """
  1595. This function determines the title level based on the font size of the title.
  1596. Parameters
  1597. ----------
  1598. title_blocks : list
  1599. Returns
  1600. -------
  1601. title_blocks : list
  1602. """
  1603. font_sizes = np.array([safe_get(tb["block"], "block_font_size", 0) for tb in title_blocks])
  1604. # Use the mean and std of font sizes to remove extreme values
  1605. mean_font_size = np.mean(font_sizes)
  1606. std_font_size = np.std(font_sizes)
  1607. min_extreme_font_size = mean_font_size - std_font_size # type: ignore
  1608. max_extreme_font_size = mean_font_size + std_font_size # type: ignore
  1609. # Compute the threshold for title level
  1610. middle_font_sizes = font_sizes[(font_sizes > min_extreme_font_size) & (font_sizes < max_extreme_font_size)]
  1611. if middle_font_sizes.size > 0:
  1612. middle_mean_font_size = np.mean(middle_font_sizes)
  1613. level_threshold = middle_mean_font_size
  1614. else:
  1615. level_threshold = mean_font_size
  1616. for tb in title_blocks:
  1617. title_block = tb["block"]
  1618. title_font_size = safe_get(title_block, "block_font_size", 0)
  1619. current_level = 1 # Initialize title level, the biggest level is 1
  1620. # print(f"Before adjustment by font size, {current_level}")
  1621. if title_font_size >= max_extreme_font_size:
  1622. current_level = 1
  1623. elif title_font_size <= min_extreme_font_size:
  1624. current_level = 3
  1625. elif float(title_font_size) >= float(level_threshold):
  1626. current_level = 2
  1627. else:
  1628. current_level = 3
  1629. # print(f"After adjustment by font size, {current_level}")
  1630. title_block["block_title_level"] = current_level
  1631. return title_blocks
  1632. def batch_recog_title_level(self, pdf_dic):
  1633. """
  1634. This function batch process the blocks to recognize title level.
  1635. Parameters
  1636. ----------
  1637. pdf_dict : dict
  1638. result dictionary
  1639. Returns
  1640. -------
  1641. pdf_dict : dict
  1642. result dictionary
  1643. """
  1644. title_blocks = []
  1645. # Collect all titles
  1646. for page_id, blocks in pdf_dic.items():
  1647. if page_id.startswith("page_"):
  1648. para_blocks = blocks.get("para_blocks", [])
  1649. for block in para_blocks:
  1650. if block.get("is_block_title"):
  1651. title_obj = {"page_id": page_id, "block": block}
  1652. title_blocks.append(title_obj)
  1653. # Determine title level
  1654. if title_blocks:
  1655. # Determine title level based on font size
  1656. title_blocks = self._recog_title_level(title_blocks)
  1657. return pdf_dic
  1658. class BlockTerminationProcessor:
  1659. """
  1660. This class is used to process the block termination.
  1661. """
  1662. def __init__(self) -> None:
  1663. pass
  1664. def _is_consistent_lines(
  1665. self,
  1666. curr_line,
  1667. prev_line,
  1668. next_line,
  1669. consistent_direction, # 0 for prev, 1 for next, 2 for both
  1670. ):
  1671. """
  1672. This function checks if the line is consistent with its neighbors
  1673. Parameters
  1674. ----------
  1675. curr_line : dict
  1676. current line
  1677. prev_line : dict
  1678. previous line
  1679. next_line : dict
  1680. next line
  1681. consistent_direction : int
  1682. 0 for prev, 1 for next, 2 for both
  1683. Returns
  1684. -------
  1685. bool
  1686. True if the line is consistent with its neighbors, False otherwise.
  1687. """
  1688. curr_line_font_size = curr_line["spans"][0]["size"]
  1689. curr_line_font_type = curr_line["spans"][0]["font"].lower()
  1690. if consistent_direction == 0:
  1691. if prev_line:
  1692. prev_line_font_size = prev_line["spans"][0]["size"]
  1693. prev_line_font_type = prev_line["spans"][0]["font"].lower()
  1694. return curr_line_font_size == prev_line_font_size and curr_line_font_type == prev_line_font_type
  1695. else:
  1696. return False
  1697. elif consistent_direction == 1:
  1698. if next_line:
  1699. next_line_font_size = next_line["spans"][0]["size"]
  1700. next_line_font_type = next_line["spans"][0]["font"].lower()
  1701. return curr_line_font_size == next_line_font_size and curr_line_font_type == next_line_font_type
  1702. else:
  1703. return False
  1704. elif consistent_direction == 2:
  1705. if prev_line and next_line:
  1706. prev_line_font_size = prev_line["spans"][0]["size"]
  1707. prev_line_font_type = prev_line["spans"][0]["font"].lower()
  1708. next_line_font_size = next_line["spans"][0]["size"]
  1709. next_line_font_type = next_line["spans"][0]["font"].lower()
  1710. return (curr_line_font_size == prev_line_font_size and curr_line_font_type == prev_line_font_type) and (
  1711. curr_line_font_size == next_line_font_size and curr_line_font_type == next_line_font_type
  1712. )
  1713. else:
  1714. return False
  1715. else:
  1716. return False
  1717. def _is_regular_line(self, curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, X0, X1, avg_line_height):
  1718. """
  1719. This function checks if the line is a regular line
  1720. Parameters
  1721. ----------
  1722. curr_line_bbox : list
  1723. bbox of the current line
  1724. prev_line_bbox : list
  1725. bbox of the previous line
  1726. next_line_bbox : list
  1727. bbox of the next line
  1728. avg_char_width : float
  1729. average of char widths
  1730. X0 : float
  1731. median of x0 values, which represents the left average boundary of the page
  1732. X1 : float
  1733. median of x1 values, which represents the right average boundary of the page
  1734. avg_line_height : float
  1735. average of line heights
  1736. Returns
  1737. -------
  1738. bool
  1739. True if the line is a regular line, False otherwise.
  1740. """
  1741. horizontal_ratio = 0.5
  1742. vertical_ratio = 0.5
  1743. horizontal_thres = horizontal_ratio * avg_char_width
  1744. vertical_thres = vertical_ratio * avg_line_height
  1745. x0, y0, x1, y1 = curr_line_bbox
  1746. x0_near_X0 = abs(x0 - X0) < horizontal_thres
  1747. x1_near_X1 = abs(x1 - X1) < horizontal_thres
  1748. prev_line_is_end_of_para = prev_line_bbox and (abs(prev_line_bbox[2] - X1) > avg_char_width)
  1749. sufficient_spacing_above = False
  1750. if prev_line_bbox:
  1751. vertical_spacing_above = y1 - prev_line_bbox[3]
  1752. sufficient_spacing_above = vertical_spacing_above > vertical_thres
  1753. sufficient_spacing_below = False
  1754. if next_line_bbox:
  1755. vertical_spacing_below = next_line_bbox[1] - y0
  1756. sufficient_spacing_below = vertical_spacing_below > vertical_thres
  1757. return (
  1758. (sufficient_spacing_above or sufficient_spacing_below)
  1759. or (not x0_near_X0 and not x1_near_X1)
  1760. or prev_line_is_end_of_para
  1761. )
  1762. def _is_possible_start_of_para(self, curr_line, prev_line, next_line, X0, X1, avg_char_width, avg_font_size):
  1763. """
  1764. This function checks if the line is a possible start of a paragraph
  1765. Parameters
  1766. ----------
  1767. curr_line : dict
  1768. current line
  1769. prev_line : dict
  1770. previous line
  1771. next_line : dict
  1772. next line
  1773. X0 : float
  1774. median of x0 values, which represents the left average boundary of the page
  1775. X1 : float
  1776. median of x1 values, which represents the right average boundary of the page
  1777. avg_char_width : float
  1778. average of char widths
  1779. avg_line_height : float
  1780. average of line heights
  1781. Returns
  1782. -------
  1783. bool
  1784. True if the line is a possible start of a paragraph, False otherwise.
  1785. """
  1786. start_confidence = 0.5 # Initial confidence of the line being a start of a paragraph
  1787. decision_path = [] # Record the decision path
  1788. curr_line_bbox = curr_line["bbox"]
  1789. prev_line_bbox = prev_line["bbox"] if prev_line else None
  1790. next_line_bbox = next_line["bbox"] if next_line else None
  1791. indent_ratio = 1
  1792. vertical_ratio = 1.5
  1793. vertical_thres = vertical_ratio * avg_font_size
  1794. left_horizontal_ratio = 0.5
  1795. left_horizontal_thres = left_horizontal_ratio * avg_char_width
  1796. right_horizontal_ratio = 2.5
  1797. right_horizontal_thres = right_horizontal_ratio * avg_char_width
  1798. x0, y0, x1, y1 = curr_line_bbox
  1799. indent_condition = x0 > X0 + indent_ratio * avg_char_width
  1800. if indent_condition:
  1801. start_confidence += 0.2
  1802. decision_path.append("indent_condition_met")
  1803. x0_near_X0 = abs(x0 - X0) < left_horizontal_thres
  1804. if x0_near_X0:
  1805. start_confidence += 0.1
  1806. decision_path.append("x0_near_X0")
  1807. x1_near_X1 = abs(x1 - X1) < right_horizontal_thres
  1808. if x1_near_X1:
  1809. start_confidence += 0.1
  1810. decision_path.append("x1_near_X1")
  1811. if prev_line is None:
  1812. prev_line_is_end_of_para = True
  1813. start_confidence += 0.2
  1814. decision_path.append("no_prev_line")
  1815. else:
  1816. prev_line_is_end_of_para, _, _ = self._is_possible_end_of_para(prev_line, next_line, X0, X1, avg_char_width)
  1817. if prev_line_is_end_of_para:
  1818. start_confidence += 0.1
  1819. decision_path.append("prev_line_is_end_of_para")
  1820. sufficient_spacing_above = False
  1821. if prev_line_bbox:
  1822. vertical_spacing_above = y1 - prev_line_bbox[3]
  1823. sufficient_spacing_above = vertical_spacing_above > vertical_thres
  1824. if sufficient_spacing_above:
  1825. start_confidence += 0.2
  1826. decision_path.append("sufficient_spacing_above")
  1827. sufficient_spacing_below = False
  1828. if next_line_bbox:
  1829. vertical_spacing_below = next_line_bbox[1] - y0
  1830. sufficient_spacing_below = vertical_spacing_below > vertical_thres
  1831. if sufficient_spacing_below:
  1832. start_confidence += 0.2
  1833. decision_path.append("sufficient_spacing_below")
  1834. is_regular_line = self._is_regular_line(
  1835. curr_line_bbox, prev_line_bbox, next_line_bbox, avg_char_width, X0, X1, avg_font_size
  1836. )
  1837. if is_regular_line:
  1838. start_confidence += 0.1
  1839. decision_path.append("is_regular_line")
  1840. is_start_of_para = (
  1841. (sufficient_spacing_above or sufficient_spacing_below)
  1842. or (indent_condition)
  1843. or (not indent_condition and x0_near_X0 and x1_near_X1 and not is_regular_line)
  1844. or prev_line_is_end_of_para
  1845. )
  1846. return (is_start_of_para, start_confidence, decision_path)
  1847. def _is_possible_end_of_para(self, curr_line, next_line, X0, X1, avg_char_width):
  1848. """
  1849. This function checks if the line is a possible end of a paragraph
  1850. Parameters
  1851. ----------
  1852. curr_line : dict
  1853. current line
  1854. next_line : dict
  1855. next line
  1856. X0 : float
  1857. median of x0 values, which represents the left average boundary of the page
  1858. X1 : float
  1859. median of x1 values, which represents the right average boundary of the page
  1860. avg_char_width : float
  1861. average of char widths
  1862. Returns
  1863. -------
  1864. bool
  1865. True if the line is a possible end of a paragraph, False otherwise.
  1866. """
  1867. end_confidence = 0.5 # Initial confidence of the line being a end of a paragraph
  1868. decision_path = [] # Record the decision path
  1869. curr_line_bbox = curr_line["bbox"]
  1870. next_line_bbox = next_line["bbox"] if next_line else None
  1871. left_horizontal_ratio = 0.5
  1872. right_horizontal_ratio = 0.5
  1873. x0, _, x1, y1 = curr_line_bbox
  1874. next_x0, next_y0, _, _ = next_line_bbox if next_line_bbox else (0, 0, 0, 0)
  1875. x0_near_X0 = abs(x0 - X0) < left_horizontal_ratio * avg_char_width
  1876. if x0_near_X0:
  1877. end_confidence += 0.1
  1878. decision_path.append("x0_near_X0")
  1879. x1_smaller_than_X1 = x1 < X1 - right_horizontal_ratio * avg_char_width
  1880. if x1_smaller_than_X1:
  1881. end_confidence += 0.1
  1882. decision_path.append("x1_smaller_than_X1")
  1883. next_line_is_start_of_para = (
  1884. next_line_bbox
  1885. and (next_x0 > X0 + left_horizontal_ratio * avg_char_width)
  1886. and (not is_line_left_aligned_from_neighbors(curr_line_bbox, None, next_line_bbox, avg_char_width, direction=1))
  1887. )
  1888. if next_line_is_start_of_para:
  1889. end_confidence += 0.2
  1890. decision_path.append("next_line_is_start_of_para")
  1891. is_line_left_aligned_from_neighbors_bool = is_line_left_aligned_from_neighbors(
  1892. curr_line_bbox, None, next_line_bbox, avg_char_width
  1893. )
  1894. if is_line_left_aligned_from_neighbors_bool:
  1895. end_confidence += 0.1
  1896. decision_path.append("line_is_left_aligned_from_neighbors")
  1897. is_line_right_aligned_from_neighbors_bool = is_line_right_aligned_from_neighbors(
  1898. curr_line_bbox, None, next_line_bbox, avg_char_width
  1899. )
  1900. if not is_line_right_aligned_from_neighbors_bool:
  1901. end_confidence += 0.1
  1902. decision_path.append("line_is_not_right_aligned_from_neighbors")
  1903. is_end_of_para = end_with_punctuation(curr_line["text"]) and (
  1904. (x0_near_X0 and x1_smaller_than_X1)
  1905. or (is_line_left_aligned_from_neighbors_bool and not is_line_right_aligned_from_neighbors_bool)
  1906. )
  1907. return (is_end_of_para, end_confidence, decision_path)
  1908. def _cut_paras_per_block(
  1909. self,
  1910. block,
  1911. ):
  1912. """
  1913. Processes a raw block from PyMuPDF and returns the processed block.
  1914. Parameters
  1915. ----------
  1916. raw_block : dict
  1917. A raw block from pymupdf.
  1918. Returns
  1919. -------
  1920. processed_block : dict
  1921. """
  1922. def _construct_para(lines, is_block_title, para_title_level):
  1923. """
  1924. Construct a paragraph from given lines.
  1925. """
  1926. font_sizes = [span["size"] for line in lines for span in line["spans"]]
  1927. avg_font_size = sum(font_sizes) / len(font_sizes) if font_sizes else 0
  1928. font_colors = [span["color"] for line in lines for span in line["spans"]]
  1929. most_common_font_color = max(set(font_colors), key=font_colors.count) if font_colors else None
  1930. font_type_lengths = {}
  1931. for line in lines:
  1932. for span in line["spans"]:
  1933. font_type = span["font"]
  1934. bbox_width = span["bbox"][2] - span["bbox"][0]
  1935. if font_type in font_type_lengths:
  1936. font_type_lengths[font_type] += bbox_width
  1937. else:
  1938. font_type_lengths[font_type] = bbox_width
  1939. # get the font type with the longest bbox width
  1940. most_common_font_type = max(font_type_lengths, key=font_type_lengths.get) if font_type_lengths else None # type: ignore
  1941. para_bbox = calculate_para_bbox(lines)
  1942. para_text = " ".join(line["text"] for line in lines)
  1943. return {
  1944. "para_bbox": para_bbox,
  1945. "para_text": para_text,
  1946. "para_font_type": most_common_font_type,
  1947. "para_font_size": avg_font_size,
  1948. "para_font_color": most_common_font_color,
  1949. "is_para_title": is_block_title,
  1950. "para_title_level": para_title_level,
  1951. }
  1952. block_bbox = block["bbox"]
  1953. block_text = block["text"]
  1954. block_lines = block["lines"]
  1955. X0 = safe_get(block, "X0", 0)
  1956. X1 = safe_get(block, "X1", 0)
  1957. avg_char_width = safe_get(block, "avg_char_width", 0)
  1958. avg_char_height = safe_get(block, "avg_char_height", 0)
  1959. avg_font_size = safe_get(block, "avg_font_size", 0)
  1960. is_block_title = safe_get(block, "is_block_title", False)
  1961. para_title_level = safe_get(block, "block_title_level", 0)
  1962. # Segment into paragraphs
  1963. para_ranges = []
  1964. in_paragraph = False
  1965. start_idx_of_para = None
  1966. # Create the processed paragraphs
  1967. processed_paras = {}
  1968. para_bboxes = []
  1969. end_idx_of_para = 0
  1970. for line_index, line in enumerate(block_lines):
  1971. curr_line = line
  1972. prev_line = block_lines[line_index - 1] if line_index > 0 else None
  1973. next_line = block_lines[line_index + 1] if line_index < len(block_lines) - 1 else None
  1974. """
  1975. Start processing paragraphs.
  1976. """
  1977. # Check if the line is the start of a paragraph
  1978. is_start_of_para, start_confidence, decision_path = self._is_possible_start_of_para(
  1979. curr_line, prev_line, next_line, X0, X1, avg_char_width, avg_font_size
  1980. )
  1981. if not in_paragraph and is_start_of_para:
  1982. in_paragraph = True
  1983. start_idx_of_para = line_index
  1984. # print_green(">>> Start of a paragraph")
  1985. # print(" curr_line_text: ", curr_line["text"])
  1986. # print(" start_confidence: ", start_confidence)
  1987. # print(" decision_path: ", decision_path)
  1988. # Check if the line is the end of a paragraph
  1989. is_end_of_para, end_confidence, decision_path = self._is_possible_end_of_para(
  1990. curr_line, next_line, X0, X1, avg_char_width
  1991. )
  1992. if in_paragraph and (is_end_of_para or not next_line):
  1993. para_ranges.append((start_idx_of_para, line_index))
  1994. start_idx_of_para = None
  1995. in_paragraph = False
  1996. # print_red(">>> End of a paragraph")
  1997. # print(" curr_line_text: ", curr_line["text"])
  1998. # print(" end_confidence: ", end_confidence)
  1999. # print(" decision_path: ", decision_path)
  2000. # Add the last paragraph if it is not added
  2001. if in_paragraph and start_idx_of_para is not None:
  2002. para_ranges.append((start_idx_of_para, len(block_lines) - 1))
  2003. # Process the matched paragraphs
  2004. for para_index, (start_idx, end_idx) in enumerate(para_ranges):
  2005. matched_lines = block_lines[start_idx : end_idx + 1]
  2006. para_properties = _construct_para(matched_lines, is_block_title, para_title_level)
  2007. para_key = f"para_{len(processed_paras)}"
  2008. processed_paras[para_key] = para_properties
  2009. para_bboxes.append(para_properties["para_bbox"])
  2010. end_idx_of_para = end_idx + 1
  2011. # Deal with the remaining lines
  2012. if end_idx_of_para < len(block_lines):
  2013. unmatched_lines = block_lines[end_idx_of_para:]
  2014. unmatched_properties = _construct_para(unmatched_lines, is_block_title, para_title_level)
  2015. unmatched_key = f"para_{len(processed_paras)}"
  2016. processed_paras[unmatched_key] = unmatched_properties
  2017. para_bboxes.append(unmatched_properties["para_bbox"])
  2018. block["paras"] = processed_paras
  2019. return block
  2020. def batch_process_blocks(self, pdf_dict):
  2021. """
  2022. Parses the blocks of all pages.
  2023. Parameters
  2024. ----------
  2025. pdf_dict : dict
  2026. PDF dictionary.
  2027. filter_blocks : list
  2028. List of bounding boxes to filter.
  2029. Returns
  2030. -------
  2031. result_dict : dict
  2032. Result dictionary.
  2033. """
  2034. num_paras = 0
  2035. for page_id, page in pdf_dict.items():
  2036. if page_id.startswith("page_"):
  2037. para_blocks = []
  2038. if "para_blocks" in page.keys():
  2039. input_blocks = page["para_blocks"]
  2040. for input_block in input_blocks:
  2041. new_block = self._cut_paras_per_block(input_block)
  2042. para_blocks.append(new_block)
  2043. num_paras += len(new_block["paras"])
  2044. page["para_blocks"] = para_blocks
  2045. pdf_dict["statistics"]["num_paras"] = num_paras
  2046. return pdf_dict
  2047. class BlockContinuationProcessor:
  2048. """
  2049. This class is used to process the blocks to detect block continuations.
  2050. """
  2051. def __init__(self) -> None:
  2052. pass
  2053. def __is_similar_font_type(self, font_type_1, font_type_2, prefix_length_ratio=0.3):
  2054. """
  2055. This function checks if the two font types are similar.
  2056. Definition of similar font types: the two font types have a common prefix,
  2057. and the length of the common prefix is at least a certain ratio of the length of the shorter font type.
  2058. Parameters
  2059. ----------
  2060. font_type1 : str
  2061. font type 1
  2062. font_type2 : str
  2063. font type 2
  2064. prefix_length_ratio : float
  2065. minimum ratio of the common prefix length to the length of the shorter font type
  2066. Returns
  2067. -------
  2068. bool
  2069. True if the two font types are similar, False otherwise.
  2070. """
  2071. if isinstance(font_type_1, list):
  2072. font_type_1 = font_type_1[0] if font_type_1 else ""
  2073. if isinstance(font_type_2, list):
  2074. font_type_2 = font_type_2[0] if font_type_2 else ""
  2075. if font_type_1 == font_type_2:
  2076. return True
  2077. # Find the length of the common prefix
  2078. common_prefix_length = len(os.path.commonprefix([font_type_1, font_type_2]))
  2079. # Calculate the minimum prefix length based on the ratio
  2080. min_prefix_length = int(min(len(font_type_1), len(font_type_2)) * prefix_length_ratio)
  2081. return common_prefix_length >= min_prefix_length
  2082. def __is_same_block_font(self, block_1, block_2):
  2083. """
  2084. This function compares the font of block1 and block2
  2085. Parameters
  2086. ----------
  2087. block1 : dict
  2088. block1
  2089. block2 : dict
  2090. block2
  2091. Returns
  2092. -------
  2093. is_same : bool
  2094. True if block1 and block2 have the same font, else False
  2095. """
  2096. block_1_font_type = safe_get(block_1, "block_font_type", "")
  2097. block_1_font_size = safe_get(block_1, "block_font_size", 0)
  2098. block_1_avg_char_width = safe_get(block_1, "avg_char_width", 0)
  2099. block_2_font_type = safe_get(block_2, "block_font_type", "")
  2100. block_2_font_size = safe_get(block_2, "block_font_size", 0)
  2101. block_2_avg_char_width = safe_get(block_2, "avg_char_width", 0)
  2102. if isinstance(block_1_font_size, list):
  2103. block_1_font_size = block_1_font_size[0] if block_1_font_size else 0
  2104. if isinstance(block_2_font_size, list):
  2105. block_2_font_size = block_2_font_size[0] if block_2_font_size else 0
  2106. block_1_text = safe_get(block_1, "text", "")
  2107. block_2_text = safe_get(block_2, "text", "")
  2108. if block_1_avg_char_width == 0 or block_2_avg_char_width == 0:
  2109. return False
  2110. if not block_1_text or not block_2_text:
  2111. return False
  2112. else:
  2113. text_len_ratio = len(block_2_text) / len(block_1_text)
  2114. if text_len_ratio < 0.2:
  2115. avg_char_width_condition = (
  2116. abs(block_1_avg_char_width - block_2_avg_char_width) / min(block_1_avg_char_width, block_2_avg_char_width)
  2117. < 0.5
  2118. )
  2119. else:
  2120. avg_char_width_condition = (
  2121. abs(block_1_avg_char_width - block_2_avg_char_width) / min(block_1_avg_char_width, block_2_avg_char_width)
  2122. < 0.2
  2123. )
  2124. block_font_size_condition = abs(block_1_font_size - block_2_font_size) < 1
  2125. return (
  2126. self.__is_similar_font_type(block_1_font_type, block_2_font_type)
  2127. and avg_char_width_condition
  2128. and block_font_size_condition
  2129. )
  2130. def _is_alphabet_char(self, char):
  2131. if (char >= "\u0041" and char <= "\u005a") or (char >= "\u0061" and char <= "\u007a"):
  2132. return True
  2133. else:
  2134. return False
  2135. def _is_chinese_char(self, char):
  2136. if char >= "\u4e00" and char <= "\u9fa5":
  2137. return True
  2138. else:
  2139. return False
  2140. def _is_other_letter_char(self, char):
  2141. try:
  2142. cat = unicodedata.category(char)
  2143. if cat == "Lu" or cat == "Ll":
  2144. return not self._is_alphabet_char(char) and not self._is_chinese_char(char)
  2145. except TypeError:
  2146. print("The input to the function must be a single character.")
  2147. return False
  2148. def _is_year(self, s: str):
  2149. try:
  2150. number = int(s)
  2151. return 1900 <= number <= 2099
  2152. except ValueError:
  2153. return False
  2154. def _match_brackets(self, text):
  2155. # pattern = r"^[\(\)\[\]()【】{}{}<><>〔〕〘〙\"\'“”‘’]"
  2156. pattern = r"^[\(\)\]()】{}{}>>〕〙\"\'“”‘’]"
  2157. return bool(re.match(pattern, text))
  2158. def _is_para_font_consistent(self, para_1, para_2):
  2159. """
  2160. This function compares the font of para1 and para2
  2161. Parameters
  2162. ----------
  2163. para1 : dict
  2164. para1
  2165. para2 : dict
  2166. para2
  2167. Returns
  2168. -------
  2169. is_same : bool
  2170. True if para1 and para2 have the same font, else False
  2171. """
  2172. if para_1 is None or para_2 is None:
  2173. return False
  2174. para_1_font_type = safe_get(para_1, "para_font_type", "")
  2175. para_1_font_size = safe_get(para_1, "para_font_size", 0)
  2176. para_1_font_color = safe_get(para_1, "para_font_color", "")
  2177. para_2_font_type = safe_get(para_2, "para_font_type", "")
  2178. para_2_font_size = safe_get(para_2, "para_font_size", 0)
  2179. para_2_font_color = safe_get(para_2, "para_font_color", "")
  2180. if isinstance(para_1_font_type, list): # get the most common font type
  2181. para_1_font_type = max(set(para_1_font_type), key=para_1_font_type.count)
  2182. if isinstance(para_2_font_type, list):
  2183. para_2_font_type = max(set(para_2_font_type), key=para_2_font_type.count)
  2184. if isinstance(para_1_font_size, list): # compute average font type
  2185. para_1_font_size = sum(para_1_font_size) / len(para_1_font_size)
  2186. if isinstance(para_2_font_size, list): # compute average font type
  2187. para_2_font_size = sum(para_2_font_size) / len(para_2_font_size)
  2188. return (
  2189. self.__is_similar_font_type(para_1_font_type, para_2_font_type)
  2190. and abs(para_1_font_size - para_2_font_size) < 1.5
  2191. # and para_font_color1 == para_font_color2
  2192. )
  2193. def _is_para_puncs_consistent(self, para_1, para_2):
  2194. """
  2195. This function determines whether para1 and para2 are originally from the same paragraph by checking the puncs of para1(former) and para2(latter)
  2196. Parameters
  2197. ----------
  2198. para1 : dict
  2199. para1
  2200. para2 : dict
  2201. para2
  2202. Returns
  2203. -------
  2204. is_same : bool
  2205. True if para1 and para2 are from the same paragraph by using the puncs, else False
  2206. """
  2207. para_1_text = safe_get(para_1, "para_text", "").strip()
  2208. para_2_text = safe_get(para_2, "para_text", "").strip()
  2209. para_1_bboxes = safe_get(para_1, "para_bbox", [])
  2210. para_1_font_sizes = safe_get(para_1, "para_font_size", 0)
  2211. para_2_bboxes = safe_get(para_2, "para_bbox", [])
  2212. para_2_font_sizes = safe_get(para_2, "para_font_size", 0)
  2213. # print_yellow(" Features of determine puncs_consistent:")
  2214. # print(f" para_1_text: {para_1_text}")
  2215. # print(f" para_2_text: {para_2_text}")
  2216. # print(f" para_1_bboxes: {para_1_bboxes}")
  2217. # print(f" para_2_bboxes: {para_2_bboxes}")
  2218. # print(f" para_1_font_sizes: {para_1_font_sizes}")
  2219. # print(f" para_2_font_sizes: {para_2_font_sizes}")
  2220. if is_nested_list(para_1_bboxes):
  2221. x0_1, y0_1, x1_1, y1_1 = para_1_bboxes[-1]
  2222. else:
  2223. x0_1, y0_1, x1_1, y1_1 = para_1_bboxes
  2224. if is_nested_list(para_2_bboxes):
  2225. x0_2, y0_2, x1_2, y1_2 = para_2_bboxes[0]
  2226. para_2_font_sizes = para_2_font_sizes[0] # type: ignore
  2227. else:
  2228. x0_2, y0_2, x1_2, y1_2 = para_2_bboxes
  2229. right_align_threshold = 0.5 * (para_1_font_sizes + para_2_font_sizes) * 0.8
  2230. are_two_paras_right_aligned = abs(x1_1 - x1_2) < right_align_threshold
  2231. left_indent_threshold = 0.5 * (para_1_font_sizes + para_2_font_sizes) * 0.8
  2232. is_para1_left_indent_than_papa2 = x0_1 - x0_2 > left_indent_threshold
  2233. is_para2_left_indent_than_papa1 = x0_2 - x0_1 > left_indent_threshold
  2234. # Check if either para_text1 or para_text2 is empty
  2235. if not para_1_text or not para_2_text:
  2236. return False
  2237. # Define the end puncs for a sentence to end and hyphen
  2238. end_puncs = [".", "?", "!", "。", "?", "!", "…"]
  2239. hyphen = ["-", "—"]
  2240. # Check if para_text1 ends with either hyphen or non-end punctuation or spaces
  2241. para_1_end_with_hyphen = para_1_text and para_1_text[-1] in hyphen
  2242. para_1_end_with_end_punc = para_1_text and para_1_text[-1] in end_puncs
  2243. para_1_end_with_space = para_1_text and para_1_text[-1] == " "
  2244. para_1_not_end_with_end_punc = para_1_text and para_1_text[-1] not in end_puncs
  2245. # print_yellow(f" para_1_end_with_hyphen: {para_1_end_with_hyphen}")
  2246. # print_yellow(f" para_1_end_with_end_punc: {para_1_end_with_end_punc}")
  2247. # print_yellow(f" para_1_not_end_with_end_punc: {para_1_not_end_with_end_punc}")
  2248. # print_yellow(f" para_1_end_with_space: {para_1_end_with_space}")
  2249. if para_1_end_with_hyphen: # If para_text1 ends with hyphen
  2250. # print_red(f"para_1 is end with hyphen.")
  2251. para_2_is_consistent = para_2_text and (
  2252. para_2_text[0] in hyphen
  2253. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
  2254. or (self._is_chinese_char(para_2_text[0]))
  2255. or (self._is_other_letter_char(para_2_text[0]))
  2256. )
  2257. if para_2_is_consistent:
  2258. # print(f"para_2 is consistent.\n")
  2259. return True
  2260. else:
  2261. # print(f"para_2 is not consistent.\n")
  2262. pass
  2263. elif para_1_end_with_end_punc: # If para_text1 ends with ending punctuations
  2264. # print_red(f"para_1 is end with end_punc.")
  2265. para_2_is_consistent = (
  2266. para_2_text
  2267. and (
  2268. para_2_text[0]
  2269. == " "
  2270. # or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].isupper())
  2271. # or (self._is_chinese_char(para_2_text[0]))
  2272. # or (self._is_other_letter_char(para_2_text[0]))
  2273. )
  2274. and not is_para2_left_indent_than_papa1
  2275. )
  2276. if para_2_is_consistent:
  2277. # print(f"para_2 is consistent.\n")
  2278. return True
  2279. else:
  2280. # print(f"para_2 is not consistent.\n")
  2281. pass
  2282. elif para_1_not_end_with_end_punc: # If para_text1 is not end with ending punctuations
  2283. # print_red(f"para_1 is NOT end with end_punc.")
  2284. para_2_is_consistent = para_2_text and (
  2285. para_2_text[0] == " "
  2286. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
  2287. or (self._is_alphabet_char(para_2_text[0]))
  2288. or (self._is_year(para_2_text[0:4]))
  2289. or (are_two_paras_right_aligned or is_para1_left_indent_than_papa2)
  2290. or (self._is_chinese_char(para_2_text[0]))
  2291. or (self._is_other_letter_char(para_2_text[0]))
  2292. or (self._match_brackets(para_2_text[0]))
  2293. )
  2294. if para_2_is_consistent:
  2295. # print(f"para_2 is consistent.\n")
  2296. return True
  2297. else:
  2298. # print(f"para_2 is not consistent.\n")
  2299. pass
  2300. elif para_1_end_with_space: # If para_text1 ends with space
  2301. # print_red(f"para_1 is end with space.")
  2302. para_2_is_consistent = para_2_text and (
  2303. para_2_text[0] == " "
  2304. or (self._is_alphabet_char(para_2_text[0]) and para_2_text[0].islower())
  2305. or (self._is_chinese_char(para_2_text[0]))
  2306. or (self._is_other_letter_char(para_2_text[0]))
  2307. )
  2308. if para_2_is_consistent:
  2309. # print(f"para_2 is consistent.\n")
  2310. return True
  2311. else:
  2312. pass
  2313. # print(f"para_2 is not consistent.\n")
  2314. return False
  2315. def _is_block_consistent(self, block_1, block_2):
  2316. """
  2317. This function determines whether block1 and block2 are originally from the same block
  2318. Parameters
  2319. ----------
  2320. block1 : dict
  2321. block1s
  2322. block2 : dict
  2323. block2
  2324. Returns
  2325. -------
  2326. is_same : bool
  2327. True if block1 and block2 are from the same block, else False
  2328. """
  2329. return self.__is_same_block_font(block_1, block_2)
  2330. def _is_para_continued(self, para_1, para_2):
  2331. """
  2332. This function determines whether para1 and para2 are originally from the same paragraph
  2333. Parameters
  2334. ----------
  2335. para1 : dict
  2336. para1
  2337. para2 : dict
  2338. para2
  2339. Returns
  2340. -------
  2341. is_same : bool
  2342. True if para1 and para2 are from the same paragraph, else False
  2343. """
  2344. is_para_font_consistent = self._is_para_font_consistent(para_1, para_2)
  2345. is_para_puncs_consistent = self._is_para_puncs_consistent(para_1, para_2)
  2346. return is_para_font_consistent and is_para_puncs_consistent
  2347. def _are_boundaries_of_block_consistent(self, block_1, block_2):
  2348. """
  2349. This function checks if the boundaries of block1 and block2 are consistent
  2350. Parameters
  2351. ----------
  2352. block1 : dict
  2353. block1
  2354. block2 : dict
  2355. block2
  2356. Returns
  2357. -------
  2358. is_consistent : bool
  2359. True if the boundaries of block1 and block2 are consistent, else False
  2360. """
  2361. last_line_of_block_1 = block_1["lines"][-1]
  2362. first_line_of_block_2 = block_2["lines"][0]
  2363. spans_of_last_line_of_block_1 = last_line_of_block_1["spans"]
  2364. spans_of_first_line_of_block_2 = first_line_of_block_2["spans"]
  2365. font_type_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["font"].lower()
  2366. font_size_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["size"]
  2367. font_color_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["color"]
  2368. font_flags_of_last_line_of_block_1 = spans_of_last_line_of_block_1[0]["flags"]
  2369. font_type_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["font"].lower()
  2370. font_size_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["size"]
  2371. font_color_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["color"]
  2372. font_flags_of_first_line_of_block_2 = spans_of_first_line_of_block_2[0]["flags"]
  2373. return (
  2374. self.__is_similar_font_type(font_type_of_last_line_of_block_1, font_type_of_first_line_of_block_2)
  2375. and abs(font_size_of_last_line_of_block_1 - font_size_of_first_line_of_block_2) < 1
  2376. # and font_color_of_last_line_of_block1 == font_color_of_first_line_of_block2
  2377. and font_flags_of_last_line_of_block_1 == font_flags_of_first_line_of_block_2
  2378. )
  2379. def should_merge_next_para(self, curr_para, next_para):
  2380. """
  2381. This function checks if the next_para should be merged into the curr_para.
  2382. Parameters
  2383. ----------
  2384. curr_para : dict
  2385. The current paragraph.
  2386. next_para : dict
  2387. The next paragraph.
  2388. Returns
  2389. -------
  2390. bool
  2391. True if the next_para should be merged into the curr_para, False otherwise.
  2392. """
  2393. if self._is_para_continued(curr_para, next_para):
  2394. return True
  2395. else:
  2396. return False
  2397. def batch_tag_paras(self, pdf_dict):
  2398. """
  2399. This function tags the paragraphs in the pdf_dict.
  2400. Parameters
  2401. ----------
  2402. pdf_dict : dict
  2403. PDF dictionary.
  2404. Returns
  2405. -------
  2406. pdf_dict : dict
  2407. PDF dictionary with tagged paragraphs.
  2408. """
  2409. the_last_page_id = len(pdf_dict) - 1
  2410. for curr_page_idx, (curr_page_id, curr_page_content) in enumerate(pdf_dict.items()):
  2411. if curr_page_id.startswith("page_") and curr_page_content.get("para_blocks", []):
  2412. para_blocks_of_curr_page = curr_page_content["para_blocks"]
  2413. next_page_idx = curr_page_idx + 1
  2414. next_page_id = f"page_{next_page_idx}"
  2415. next_page_content = pdf_dict.get(next_page_id, {})
  2416. for i, current_block in enumerate(para_blocks_of_curr_page):
  2417. for para_id, curr_para in current_block["paras"].items():
  2418. curr_para["curr_para_location"] = [
  2419. curr_page_idx,
  2420. current_block["block_id"],
  2421. int(para_id.split("_")[-1]),
  2422. ]
  2423. curr_para["next_para_location"] = None # 默认设置为None
  2424. curr_para["merge_next_para"] = False # 默认设置为False
  2425. next_block = para_blocks_of_curr_page[i + 1] if i < len(para_blocks_of_curr_page) - 1 else None
  2426. if next_block:
  2427. curr_block_last_para_key = list(current_block["paras"].keys())[-1]
  2428. curr_blk_last_para = current_block["paras"][curr_block_last_para_key]
  2429. next_block_first_para_key = list(next_block["paras"].keys())[0]
  2430. next_blk_first_para = next_block["paras"][next_block_first_para_key]
  2431. if self.should_merge_next_para(curr_blk_last_para, next_blk_first_para):
  2432. curr_blk_last_para["next_para_location"] = [
  2433. curr_page_idx,
  2434. next_block["block_id"],
  2435. int(next_block_first_para_key.split("_")[-1]),
  2436. ]
  2437. curr_blk_last_para["merge_next_para"] = True
  2438. else:
  2439. # Handle the case where the next block is in a different page
  2440. curr_block_last_para_key = list(current_block["paras"].keys())[-1]
  2441. curr_blk_last_para = current_block["paras"][curr_block_last_para_key]
  2442. while not next_page_content.get("para_blocks", []) and next_page_idx <= the_last_page_id:
  2443. next_page_idx += 1
  2444. next_page_id = f"page_{next_page_idx}"
  2445. next_page_content = pdf_dict.get(next_page_id, {})
  2446. if next_page_content.get("para_blocks", []):
  2447. next_blk_first_para_key = list(next_page_content["para_blocks"][0]["paras"].keys())[0]
  2448. next_blk_first_para = next_page_content["para_blocks"][0]["paras"][next_blk_first_para_key]
  2449. if self.should_merge_next_para(curr_blk_last_para, next_blk_first_para):
  2450. curr_blk_last_para["next_para_location"] = [
  2451. next_page_idx,
  2452. next_page_content["para_blocks"][0]["block_id"],
  2453. int(next_blk_first_para_key.split("_")[-1]),
  2454. ]
  2455. curr_blk_last_para["merge_next_para"] = True
  2456. return pdf_dict
  2457. def find_block_by_id(self, para_blocks, block_id):
  2458. """
  2459. This function finds a block by its id.
  2460. Parameters
  2461. ----------
  2462. para_blocks : list
  2463. List of blocks.
  2464. block_id : int
  2465. Id of the block to find.
  2466. Returns
  2467. -------
  2468. block : dict
  2469. The block with the given id.
  2470. """
  2471. for block in para_blocks:
  2472. if block.get("block_id") == block_id:
  2473. return block
  2474. return None
  2475. def batch_merge_paras(self, pdf_dict):
  2476. """
  2477. This function merges the paragraphs in the pdf_dict.
  2478. Parameters
  2479. ----------
  2480. pdf_dict : dict
  2481. PDF dictionary.
  2482. Returns
  2483. -------
  2484. pdf_dict : dict
  2485. PDF dictionary with merged paragraphs.
  2486. """
  2487. for page_id, page_content in pdf_dict.items():
  2488. if page_id.startswith("page_") and page_content.get("para_blocks", []):
  2489. para_blocks_of_page = page_content["para_blocks"]
  2490. for i in range(len(para_blocks_of_page)):
  2491. current_block = para_blocks_of_page[i]
  2492. paras = current_block["paras"]
  2493. for para_id, curr_para in list(paras.items()):
  2494. # 跳过标题段落
  2495. if curr_para.get("is_para_title"):
  2496. continue
  2497. while curr_para.get("merge_next_para"):
  2498. next_para_location = curr_para.get("next_para_location")
  2499. if not next_para_location:
  2500. break
  2501. next_page_idx, next_block_id, next_para_id = next_para_location
  2502. next_page_id = f"page_{next_page_idx}"
  2503. next_page_content = pdf_dict.get(next_page_id)
  2504. if not next_page_content:
  2505. break
  2506. next_block = self.find_block_by_id(next_page_content.get("para_blocks", []), next_block_id)
  2507. if not next_block:
  2508. break
  2509. next_para = next_block["paras"].get(f"para_{next_para_id}")
  2510. if not next_para or next_para.get("is_para_title"):
  2511. break
  2512. # 合并段落文本
  2513. curr_para_text = curr_para.get("para_text", "")
  2514. next_para_text = next_para.get("para_text", "")
  2515. curr_para["para_text"] = curr_para_text + " " + next_para_text
  2516. # 更新 next_para_location
  2517. curr_para["next_para_location"] = next_para.get("next_para_location")
  2518. # 将下一个段落文本置为空,表示已被合并
  2519. next_para["para_text"] = ""
  2520. # 更新 merge_next_para 标记
  2521. curr_para["merge_next_para"] = next_para.get("merge_next_para", False)
  2522. return pdf_dict
  2523. class DrawAnnos:
  2524. """
  2525. This class draws annotations on the pdf file
  2526. ----------------------------------------
  2527. Color Code
  2528. ----------------------------------------
  2529. Red: (1, 0, 0)
  2530. Green: (0, 1, 0)
  2531. Blue: (0, 0, 1)
  2532. Yellow: (1, 1, 0) - mix of red and green
  2533. Cyan: (0, 1, 1) - mix of green and blue
  2534. Magenta: (1, 0, 1) - mix of red and blue
  2535. White: (1, 1, 1) - red, green and blue full intensity
  2536. Black: (0, 0, 0) - no color component whatsoever
  2537. Gray: (0.5, 0.5, 0.5) - equal and medium intensity of red, green and blue color components
  2538. Orange: (1, 0.65, 0) - maximum intensity of red, medium intensity of green, no blue component
  2539. """
  2540. def __init__(self) -> None:
  2541. pass
  2542. def __is_nested_list(self, lst):
  2543. """
  2544. This function returns True if the given list is a nested list of any degree.
  2545. """
  2546. if isinstance(lst, list):
  2547. return any(self.__is_nested_list(i) for i in lst) or any(isinstance(i, list) for i in lst)
  2548. return False
  2549. def __valid_rect(self, bbox):
  2550. # Ensure that the rectangle is not empty or invalid
  2551. if isinstance(bbox[0], list):
  2552. return False # It's a nested list, hence it can't be valid rect
  2553. else:
  2554. return bbox[0] < bbox[2] and bbox[1] < bbox[3]
  2555. def __draw_nested_boxes(self, page, nested_bbox, color=(0, 1, 1)):
  2556. """
  2557. This function draws the nested boxes
  2558. Parameters
  2559. ----------
  2560. page : fitz.Page
  2561. page
  2562. nested_bbox : list
  2563. nested bbox
  2564. color : tuple
  2565. color, by default (0, 1, 1) # draw with cyan color for combined paragraph
  2566. """
  2567. if self.__is_nested_list(nested_bbox): # If it's a nested list
  2568. for bbox in nested_bbox:
  2569. self.__draw_nested_boxes(page, bbox, color) # Recursively call the function
  2570. elif self.__valid_rect(nested_bbox): # If valid rectangle
  2571. para_rect = fitz.Rect(nested_bbox)
  2572. para_anno = page.add_rect_annot(para_rect)
  2573. para_anno.set_colors(stroke=color) # draw with cyan color for combined paragraph
  2574. para_anno.set_border(width=1)
  2575. para_anno.update()
  2576. def draw_annos(self, input_pdf_path, pdf_dic, output_pdf_path):
  2577. """
  2578. This function draws annotations on the pdf file.
  2579. Parameters
  2580. ----------
  2581. input_pdf_path : str
  2582. path to the input pdf file
  2583. pdf_dic : dict
  2584. pdf dictionary
  2585. output_pdf_path : str
  2586. path to the output pdf file
  2587. pdf_dic : dict
  2588. pdf dictionary
  2589. """
  2590. pdf_doc = open_pdf(input_pdf_path)
  2591. if pdf_dic is None:
  2592. pdf_dic = {}
  2593. if output_pdf_path is None:
  2594. output_pdf_path = input_pdf_path.replace(".pdf", "_anno.pdf")
  2595. for page_id, page in enumerate(pdf_doc): # type: ignore
  2596. page_key = f"page_{page_id}"
  2597. for ele_key, ele_data in pdf_dic[page_key].items():
  2598. if ele_key == "para_blocks":
  2599. para_blocks = ele_data
  2600. for para_block in para_blocks:
  2601. if "paras" in para_block.keys():
  2602. paras = para_block["paras"]
  2603. for para_key, para_content in paras.items():
  2604. para_bbox = para_content["para_bbox"]
  2605. # print(f"para_bbox: {para_bbox}")
  2606. # print(f"is a nested list: {self.__is_nested_list(para_bbox)}")
  2607. if self.__is_nested_list(para_bbox) and len(para_bbox) > 1:
  2608. color = (0, 1, 1)
  2609. self.__draw_nested_boxes(
  2610. page, para_bbox, color
  2611. ) # draw with cyan color for combined paragraph
  2612. else:
  2613. if self.__valid_rect(para_bbox):
  2614. para_rect = fitz.Rect(para_bbox)
  2615. para_anno = page.add_rect_annot(para_rect)
  2616. para_anno.set_colors(stroke=(0, 1, 0)) # draw with green color for normal paragraph
  2617. para_anno.set_border(width=0.5)
  2618. para_anno.update()
  2619. is_para_title = para_content["is_para_title"]
  2620. if is_para_title:
  2621. if self.__is_nested_list(para_content["para_bbox"]) and len(para_content["para_bbox"]) > 1:
  2622. color = (0, 0, 1)
  2623. self.__draw_nested_boxes(
  2624. page, para_content["para_bbox"], color
  2625. ) # draw with cyan color for combined title
  2626. else:
  2627. if self.__valid_rect(para_content["para_bbox"]):
  2628. para_rect = fitz.Rect(para_content["para_bbox"])
  2629. if self.__valid_rect(para_content["para_bbox"]):
  2630. para_anno = page.add_rect_annot(para_rect)
  2631. para_anno.set_colors(stroke=(0, 0, 1)) # draw with blue color for normal title
  2632. para_anno.set_border(width=0.5)
  2633. para_anno.update()
  2634. pdf_doc.save(output_pdf_path)
  2635. pdf_doc.close()
  2636. class ParaProcessPipeline:
  2637. def __init__(self) -> None:
  2638. pass
  2639. def para_process_pipeline(self, pdf_info_dict, para_debug_mode=None, input_pdf_path=None, output_pdf_path=None):
  2640. """
  2641. This function processes the paragraphs, including:
  2642. 1. Read raw input json file into pdf_dic
  2643. 2. Detect and replace equations
  2644. 3. Combine spans into a natural line
  2645. 4. Check if the paragraphs are inside bboxes passed from "layout_bboxes" key
  2646. 5. Compute statistics for each block
  2647. 6. Detect titles in the document
  2648. 7. Detect paragraphs inside each block
  2649. 8. Divide the level of the titles
  2650. 9. Detect and combine paragraphs from different blocks into one paragraph
  2651. 10. Check whether the final results after checking headings, dividing paragraphs within blocks, and merging paragraphs between blocks are plausible and reasonable.
  2652. 11. Draw annotations on the pdf file
  2653. Parameters
  2654. ----------
  2655. pdf_dic_json_fpath : str
  2656. path to the pdf dictionary json file.
  2657. Notice: data noises, including overlap blocks, header, footer, watermark, vertical margin note have been removed already.
  2658. input_pdf_doc : str
  2659. path to the input pdf file
  2660. output_pdf_path : str
  2661. path to the output pdf file
  2662. Returns
  2663. -------
  2664. pdf_dict : dict
  2665. result dictionary
  2666. """
  2667. error_info = None
  2668. output_json_file = ""
  2669. output_dir = ""
  2670. if input_pdf_path is not None:
  2671. input_pdf_path = os.path.abspath(input_pdf_path)
  2672. # print_green_on_red(f">>>>>>>>>>>>>>>>>>> Process the paragraphs of {input_pdf_path}")
  2673. if output_pdf_path is not None:
  2674. output_dir = os.path.dirname(output_pdf_path)
  2675. output_json_file = f"{output_dir}/pdf_dic.json"
  2676. def __save_pdf_dic(pdf_dic, output_pdf_path, stage="0", para_debug_mode=para_debug_mode):
  2677. """
  2678. Save the pdf_dic to a json file
  2679. """
  2680. output_pdf_file_name = os.path.basename(output_pdf_path)
  2681. # output_dir = os.path.dirname(output_pdf_path)
  2682. output_dir = "\\tmp\\pdf_parse"
  2683. output_pdf_file_name = output_pdf_file_name.replace(".pdf", f"_stage_{stage}.json")
  2684. pdf_dic_json_fpath = os.path.join(output_dir, output_pdf_file_name)
  2685. if not os.path.exists(output_dir):
  2686. os.makedirs(output_dir)
  2687. if para_debug_mode == "full":
  2688. with open(pdf_dic_json_fpath, "w", encoding="utf-8") as f:
  2689. json.dump(pdf_dic, f, indent=2, ensure_ascii=False)
  2690. # Validate the output already exists
  2691. if not os.path.exists(pdf_dic_json_fpath):
  2692. print_red(f"Failed to save the pdf_dic to {pdf_dic_json_fpath}")
  2693. return None
  2694. else:
  2695. print_green(f"Succeed to save the pdf_dic to {pdf_dic_json_fpath}")
  2696. return pdf_dic_json_fpath
  2697. """
  2698. Preprocess the lines of block
  2699. """
  2700. # Combine spans into a natural line
  2701. rawBlockProcessor = RawBlockProcessor()
  2702. pdf_dic = rawBlockProcessor.batch_process_blocks(pdf_info_dict)
  2703. # print(f"pdf_dic['page_0']['para_blocks'][0]: {pdf_dic['page_0']['para_blocks'][0]}", end="\n\n")
  2704. # Check if the paragraphs are inside bboxes passed from "layout_bboxes" key
  2705. layoutFilter = LayoutFilterProcessor()
  2706. pdf_dic = layoutFilter.batch_process_blocks(pdf_dic)
  2707. # Compute statistics for each block
  2708. blockStatisticsCalculator = BlockStatisticsCalculator()
  2709. pdf_dic = blockStatisticsCalculator.batch_process_blocks(pdf_dic)
  2710. # print(f"pdf_dic['page_0']['para_blocks'][0]: {pdf_dic['page_0']['para_blocks'][0]}", end="\n\n")
  2711. # Compute statistics for all blocks(namely this pdf document)
  2712. docStatisticsCalculator = DocStatisticsCalculator()
  2713. pdf_dic = docStatisticsCalculator.calc_stats_of_doc(pdf_dic)
  2714. # print(f"pdf_dic['statistics']: {pdf_dic['statistics']}", end="\n\n")
  2715. # Dump the first three stages of pdf_dic to a json file
  2716. if para_debug_mode == "full":
  2717. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="0", para_debug_mode=para_debug_mode)
  2718. """
  2719. Detect titles in the document
  2720. """
  2721. doc_statistics = pdf_dic["statistics"]
  2722. titleProcessor = TitleProcessor(doc_statistics)
  2723. titleProcessor.stage = 0
  2724. pdf_dic = titleProcessor.batch_detect_titles(pdf_dic)
  2725. if para_debug_mode == "full":
  2726. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="1", para_debug_mode=para_debug_mode)
  2727. """
  2728. Detect and divide the level of the titles
  2729. """
  2730. titleProcessor = TitleProcessor()
  2731. pdf_dic = titleProcessor.batch_recog_title_level(pdf_dic)
  2732. if para_debug_mode == "full":
  2733. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="2", para_debug_mode=para_debug_mode)
  2734. """
  2735. Detect and split paragraphs inside each block
  2736. """
  2737. blockInnerParasProcessor = BlockTerminationProcessor()
  2738. pdf_dic = blockInnerParasProcessor.batch_process_blocks(pdf_dic)
  2739. if para_debug_mode == "full":
  2740. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="3", para_debug_mode=para_debug_mode)
  2741. # pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="3", para_debug_mode="full")
  2742. # print_green(f"pdf_dic_json_fpath: {pdf_dic_json_fpath}")
  2743. """
  2744. Detect and combine paragraphs from different blocks into one paragraph
  2745. """
  2746. blockContinuationProcessor = BlockContinuationProcessor()
  2747. pdf_dic = blockContinuationProcessor.batch_tag_paras(pdf_dic)
  2748. pdf_dic = blockContinuationProcessor.batch_merge_paras(pdf_dic)
  2749. """
  2750. Detect titles in the document again
  2751. """
  2752. doc_statistics = pdf_dic["statistics"]
  2753. titleProcessor = TitleProcessor(doc_statistics)
  2754. titleProcessor.stage = 1
  2755. # pdf_dic = titleProcessor.batch_detect_titles(pdf_dic)
  2756. if para_debug_mode == "full":
  2757. pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="4", para_debug_mode=para_debug_mode)
  2758. # pdf_dic_json_fpath = __save_pdf_dic(pdf_dic, output_pdf_path, stage="4", para_debug_mode="full")
  2759. # print_green(f"pdf_dic_json_fpath: {pdf_dic_json_fpath}")
  2760. """
  2761. Discard pdf files by checking exceptions and return the error info to the caller
  2762. """
  2763. discardByException = DiscardByException()
  2764. is_discard_by_single_line_block = discardByException.discard_by_single_line_block(
  2765. pdf_dic, exception=DenseSingleLineBlockException()
  2766. )
  2767. is_discard_by_title_detection = discardByException.discard_by_title_detection(
  2768. pdf_dic, exception=TitleDetectionException()
  2769. )
  2770. is_discard_by_title_level = discardByException.discard_by_title_level(pdf_dic, exception=TitleLevelException())
  2771. is_discard_by_split_para = discardByException.discard_by_split_para(pdf_dic, exception=ParaSplitException())
  2772. is_discard_by_merge_para = discardByException.discard_by_merge_para(pdf_dic, exception=ParaMergeException())
  2773. if is_discard_by_single_line_block is not None:
  2774. error_info = is_discard_by_single_line_block
  2775. elif is_discard_by_title_detection is not None:
  2776. error_info = is_discard_by_title_detection
  2777. elif is_discard_by_title_level is not None:
  2778. error_info = is_discard_by_title_level
  2779. elif is_discard_by_split_para is not None:
  2780. error_info = is_discard_by_split_para
  2781. elif is_discard_by_merge_para is not None:
  2782. error_info = is_discard_by_merge_para
  2783. if error_info is not None:
  2784. return pdf_dic, error_info
  2785. """
  2786. Dump the final pdf_dic to a json file
  2787. """
  2788. if para_debug_mode is not None:
  2789. with open(output_json_file, "w", encoding="utf-8") as f:
  2790. json.dump(pdf_info_dict, f, ensure_ascii=False, indent=4)
  2791. """
  2792. Draw the annotations
  2793. """
  2794. if para_debug_mode is not None:
  2795. drawAnnos = DrawAnnos()
  2796. drawAnnos.draw_annos(input_pdf_path, pdf_dic, output_pdf_path)
  2797. """
  2798. Remove the intermediate files which are generated in the process of paragraph processing if debug_mode is simple
  2799. """
  2800. if para_debug_mode is not None:
  2801. for fpath in os.listdir(output_dir):
  2802. if fpath.endswith(".json") and "stage" in fpath:
  2803. os.remove(os.path.join(output_dir, fpath))
  2804. return pdf_dic, error_info
  2805. """
  2806. Run this script to test the function with Command:
  2807. python detect_para.py [pdf_path] [output_pdf_path]
  2808. Params:
  2809. - pdf_path: the path of the pdf file
  2810. - output_pdf_path: the path of the output pdf file
  2811. """
  2812. if __name__ == "__main__":
  2813. DEFAULT_PDF_PATH = (
  2814. "app/pdf_toolbox/tests/assets/paper/paper.pdf" if os.name != "nt" else "app\\pdf_toolbox\\tests\\assets\\paper\\paper.pdf"
  2815. )
  2816. input_pdf_path = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_PDF_PATH
  2817. output_pdf_path = sys.argv[2] if len(sys.argv) > 2 else input_pdf_path.split(".")[0] + "_recogPara.pdf"
  2818. output_json_path = sys.argv[3] if len(sys.argv) > 3 else input_pdf_path.split(".")[0] + "_recogPara.json"
  2819. import stat
  2820. # Remove existing output file if it exists
  2821. if os.path.exists(output_pdf_path):
  2822. os.chmod(output_pdf_path, stat.S_IWRITE)
  2823. os.remove(output_pdf_path)
  2824. input_pdf_doc = open_pdf(input_pdf_path)
  2825. # postprocess the paragraphs
  2826. paraProcessPipeline = ParaProcessPipeline()
  2827. # parse paragraph and save to json file
  2828. pdf_dic = {}
  2829. blockInnerParasProcessor = BlockTerminationProcessor()
  2830. """
  2831. Construct the pdf dictionary.
  2832. """
  2833. for page_id, page in enumerate(input_pdf_doc): # type: ignore
  2834. # print(f"Processing page {page_id}")
  2835. # print(f"page: {page}")
  2836. raw_blocks = page.get_text("dict")["blocks"]
  2837. # Save text blocks to "preproc_blocks"
  2838. preproc_blocks = []
  2839. for block in raw_blocks:
  2840. if block["type"] == 0:
  2841. preproc_blocks.append(block)
  2842. layout_bboxes = []
  2843. # Construct the pdf dictionary as schema above
  2844. page_dict = {
  2845. "para_blocks": None,
  2846. "preproc_blocks": preproc_blocks,
  2847. "images": None,
  2848. "tables": None,
  2849. "interline_equations": None,
  2850. "inline_equations": None,
  2851. "layout_bboxes": None,
  2852. "pymu_raw_blocks": None,
  2853. "global_statistic": None,
  2854. "droped_text_block": None,
  2855. "droped_image_block": None,
  2856. "droped_table_block": None,
  2857. "image_backup": None,
  2858. "table_backup": None,
  2859. }
  2860. pdf_dic[f"page_{page_id}"] = page_dict
  2861. # print(f"pdf_dic: {pdf_dic}")
  2862. with open(output_json_path, "w", encoding="utf-8") as f:
  2863. json.dump(pdf_dic, f, ensure_ascii=False, indent=4)
  2864. pdf_dic = paraProcessPipeline.para_process_pipeline(output_json_path, input_pdf_doc, output_pdf_path)